rename DATASETS_DIR to data_dir

Author: Jaan Altosaar

data.py

@@ -7,36 +7,37 @@
 import h5py
 
 
-def parse_binary_mnist():
+def parse_binary_mnist(data_dir):
   def lines_to_np_array(lines):
     return np.array([[int(i) for i in line.split()] for line in lines])
-  with open(os.path.join(DATASETS_DIR, 'binarized_mnist_train.amat')) as f:
+  with open(os.path.join(data_dir, 'binarized_mnist_train.amat')) as f:
     lines = f.readlines()
   train_data = lines_to_np_array(lines).astype('float32')
-  with open(os.path.join(DATASETS_DIR, 'binarized_mnist_valid.amat')) as f:
+  with open(os.path.join(data_dir, 'binarized_mnist_valid.amat')) as f:
     lines = f.readlines()
   validation_data = lines_to_np_array(lines).astype('float32')
-  with open(os.path.join(DATASETS_DIR, 'binarized_mnist_test.amat')) as f:
+  with open(os.path.join(data_dir, 'binarized_mnist_test.amat')) as f:
     lines = f.readlines()
   test_data = lines_to_np_array(lines).astype('float32')
   return train_data, validation_data, test_data
 
 
 def download_binary_mnist(fname):
-  DATASETS_DIR = '/tmp/'
+  data_dir = '/tmp/'
   subdatasets = ['train', 'valid', 'test']
   for subdataset in subdatasets:
     filename = 'binarized_mnist_{}.amat'.format(subdataset)
     url = 'http://www.cs.toronto.edu/~larocheh/public/datasets/binarized_mnist/binarized_mnist_{}.amat'.format(
         subdataset)
-    local_filename = os.path.join(DATASETS_DIR, filename)
+    local_filename = os.path.join(data_dir, filename)
     urllib.request.urlretrieve(url, local_filename)
 
-  train, validation, test = parse_binary_mnist()
+  train, validation, test = parse_binary_mnist(data_dir)
 
   data_dict = {'train': train, 'valid': validation, 'test': test}
   f = h5py.File(fname, 'w')
   f.create_dataset('train', data=data_dict['train'])
   f.create_dataset('valid', data=data_dict['valid'])
   f.create_dataset('test', data=data_dict['test'])
   f.close()
+  print(f'Saved binary MNIST data to: {fname}')

train_variational_autoencoder_pytorch.py

@@ -152,6 +152,7 @@ def cycle(iterable):
 def load_binary_mnist(cfg, **kwcfg):
   fname = cfg.data_dir / 'binary_mnist.h5'
   if not fname.exists():
+    print('Downloading binary MNIST data...')
     data.download_binary_mnist(fname)
   f = h5py.File(pathlib.os.path.join(pathlib.os.environ['DAT'], 'binarized_mnist.hdf5'), 'r')
   x_train = f['train'][::]