format and lint with flake8 and black

Author: Jaan Altosaar

README.md

@@ -9,39 +9,16 @@ Variational inference is used to fit the model to binarized MNIST handwritten di
 
 Blog post: https://jaan.io/what-is-variational-autoencoder-vae-tutorial/
 
-Example output with importance sampling for estimating the marginal likelihood on Hugo Larochelle's Binary MNIST dataset. Finaly marginal likelihood on the test set of `-97.10` nats.
+Example output with importance sampling for estimating the marginal likelihood on Hugo Larochelle's Binary MNIST dataset. Final marginal likelihood on the test set was `-97.10` nats after 65k iterations.
 
 ```
 $ python train_variational_autoencoder_pytorch.py --variational mean-field
-step:   0       train elbo: -558.69
-step:   0               valid elbo: -391.84     valid log p(x): -363.25
-step:   5000    train elbo: -116.09
-step:   5000            valid elbo: -112.57     valid log p(x): -107.01
-step:   10000   train elbo: -105.82
-step:   10000           valid elbo: -108.49     valid log p(x): -102.62
-step:   15000   train elbo: -106.78
-step:   15000           valid elbo: -106.97     valid log p(x): -100.97
-step:   20000   train elbo: -108.43
-step:   20000           valid elbo: -106.23     valid log p(x): -100.04
-step:   25000   train elbo: -99.68
-step:   25000           valid elbo: -104.89     valid log p(x): -98.83
-step:   30000   train elbo: -96.71
-step:   30000           valid elbo: -104.50     valid log p(x): -98.34
-step:   35000   train elbo: -98.64
-step:   35000           valid elbo: -104.05     valid log p(x): -97.87
-step:   40000   train elbo: -93.60
-step:   40000           valid elbo: -104.10     valid log p(x): -97.68
-step:   45000   train elbo: -96.45
-step:   45000           valid elbo: -104.58     valid log p(x): -97.76
-step:   50000   train elbo: -101.63
-step:   50000           valid elbo: -104.72     valid log p(x): -97.81
-step:   55000   train elbo: -106.78
-step:   55000           valid elbo: -105.14     valid log p(x): -98.06
-step:   60000   train elbo: -100.58
-step:   60000           valid elbo: -104.13     valid log p(x): -97.30
-step:   65000   train elbo: -96.19
-step:   65000           valid elbo: -104.46     valid log p(x): -97.43
-step:   65000           test elbo: -103.31      test log p(x): -97.10
+step:   0       train elbo: -558.28
+step:   0               valid elbo: -392.78     valid log p(x): -359.91
+step:   10000   train elbo: -106.67
+step:   10000           valid elbo: -109.12     valid log p(x): -103.11
+step:   20000   train elbo: -107.28
+step:   20000           valid elbo: -105.65     valid log p(x): -99.74
 ```
 
 

data.py

@@ -5,39 +5,63 @@
 import os
 import numpy as np
 import h5py
+import torch
 
 
 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(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(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(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 lines_to_np_array(lines):
+        return np.array([[int(i) for i in line.split()] for line in lines])
+
+    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(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(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):
-  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(data_dir, filename)
-    urllib.request.urlretrieve(url, local_filename)
-
-  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}')
+    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(data_dir, filename)
+        urllib.request.urlretrieve(url, local_filename)
+
+    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}")
+
+
+def load_binary_mnist(fname, batch_size, test_batch_size, use_gpu):
+    f = h5py.File(fname, "r")
+    x_train = f["train"][::]
+    x_val = f["valid"][::]
+    x_test = f["test"][::]
+    train = torch.utils.data.TensorDataset(torch.from_numpy(x_train))
+    kwargs = {"num_workers": 4, "pin_memory": True} if use_gpu else {}
+    train_loader = torch.utils.data.DataLoader(
+        train, batch_size=batch_size, shuffle=True, **kwargs
+    )
+    validation = torch.utils.data.TensorDataset(torch.from_numpy(x_val))
+    val_loader = torch.utils.data.DataLoader(
+        validation, batch_size=test_batch_size, shuffle=False, **kwargs
+    )
+    test = torch.utils.data.TensorDataset(torch.from_numpy(x_test))
+    test_loader = torch.utils.data.DataLoader(
+        test, batch_size=test_batch_size, shuffle=False, **kwargs
+    )
+    return train_loader, val_loader, test_loader