| 41 | |
| 42 | template<bool expand_labels> |
| 43 | static void setup_mnist(int *num_classes, int *num_train, int *num_test, |
| 44 | af::array &train_images, af::array &test_images, |
| 45 | af::array &train_labels, af::array &test_labels, |
| 46 | float frac) { |
| 47 | std::vector<dim_t> idims; |
| 48 | std::vector<float> idata; |
| 49 | read_idx(idims, idata, ASSETS_DIR "/examples/data/mnist/images-subset"); |
| 50 | |
| 51 | std::vector<dim_t> ldims; |
| 52 | std::vector<unsigned> ldata; |
| 53 | read_idx(ldims, ldata, ASSETS_DIR "/examples/data/mnist/labels-subset"); |
| 54 | |
| 55 | std::reverse(idims.begin(), idims.end()); |
| 56 | unsigned numdims = idims.size(); |
| 57 | af::array images = af::array(af::dim4(numdims, &idims[0]), &idata[0]); |
| 58 | |
| 59 | af::array R = af::randu(10000, 1); |
| 60 | af::array cond = R < std::min(frac, 0.8f); |
| 61 | af::array train_indices = where(cond); |
| 62 | af::array test_indices = where(!cond); |
| 63 | |
| 64 | train_images = lookup(images, train_indices, 2) / 255; |
| 65 | test_images = lookup(images, test_indices, 2) / 255; |
| 66 | |
| 67 | *num_classes = 10; |
| 68 | *num_train = train_images.dims(2); |
| 69 | *num_test = test_images.dims(2); |
| 70 | |
| 71 | if (expand_labels) { |
| 72 | train_labels = af::constant(0, *num_classes, *num_train); |
| 73 | test_labels = af::constant(0, *num_classes, *num_test); |
| 74 | |
| 75 | unsigned *h_train_idx = train_indices.host<unsigned>(); |
| 76 | unsigned *h_test_idx = test_indices.host<unsigned>(); |
| 77 | |
| 78 | for (int ii = 0; ii < *num_train; ii++) { |
| 79 | train_labels(ldata[h_train_idx[ii]], ii) = 1; |
| 80 | } |
| 81 | |
| 82 | for (int ii = 0; ii < *num_test; ii++) { |
| 83 | test_labels(ldata[h_test_idx[ii]], ii) = 1; |
| 84 | } |
| 85 | |
| 86 | af::freeHost(h_train_idx); |
| 87 | af::freeHost(h_test_idx); |
| 88 | } else { |
| 89 | af::array labels = af::array(ldims[0], &ldata[0]); |
| 90 | train_labels = labels(train_indices); |
| 91 | test_labels = labels(test_indices); |
| 92 | } |
| 93 | |
| 94 | return; |
| 95 | } |
| 96 | |
| 97 | #if 0 |
| 98 | static af::array randidx(int num, int total) |