keras mnist replaced with torchvision mnist

Author: cosmic-cortex

examples/pytorch_integration.py

@@ -6,11 +6,16 @@
 
 import torch
 import numpy as np
-from keras.datasets import mnist
+
 from torch import nn
+from torch.utils.data import DataLoader
+from torchvision.transforms import ToTensor
+from torchvision.datasets import MNIST
 from skorch import NeuralNetClassifier
+
 from modAL.models import ActiveLearner
 
+
 # build class for the skorch API
 class Torch_Model(nn.Module):
     def __init__(self,):
@@ -37,28 +42,34 @@ def forward(self, x):
         out = self.fcs(out)
         return out
 
+
 # create the classifier
+device = "cuda" if torch.cuda.is_available() else "cpu"
 classifier = NeuralNetClassifier(Torch_Model,
                                  # max_epochs=100,
                                  criterion=nn.CrossEntropyLoss,
                                  optimizer=torch.optim.Adam,
                                  train_split=None,
                                  verbose=1,
-                                 device="cuda")
+                                 device=device)
 
 """
 Data wrangling
-1. Reading data from Keras
+1. Reading data from torchvision
 2. Assembling initial training data for ActiveLearner
 3. Generating the pool
 """
 
+mnist_data = MNIST('.', download=True, transform=ToTensor())
+dataloader = DataLoader(mnist_data, shuffle=True, batch_size=60000)
+X, y = next(iter(dataloader))
+
 # read training data
-(X_train, y_train), (X_test, y_test) = mnist.load_data()
-X_train = X_train.reshape(60000, 1, 28, 28).astype('float32') / 255.
-X_test = X_test.reshape(10000, 1, 28, 28).astype('float32') / 255.
-y_train = y_train.astype('long')
-y_test  = y_test.astype('long')
+X_train, X_test, y_train, y_test = X[:50000], X[50000:], y[:50000], y[50000:]
+X_train = X_train.reshape(50000, 1, 28, 28)
+X_test = X_test.reshape(10000, 1, 28, 28)
+y_train = y_train
+y_test  = y_test
 
 # assemble initial data
 n_initial = 1000
@@ -85,7 +96,6 @@ def forward(self, x):
 n_queries = 10
 for idx in range(n_queries):
     query_idx, query_instance = learner.query(X_pool, n_instances=100)
-    print(query_idx)
     learner.teach(X_pool[query_idx], y_pool[query_idx], only_new=True)
     # remove queried instance from pool
     X_pool = np.delete(X_pool, query_idx, axis=0)