fix sample so it works on the GPU (#505)

Author: lovettchris

Samples/Tutorial Samples/PyTorch Image Classification/PyTorchTraining - Image Classification/PyTorchTraining.py

@@ -21,21 +21,21 @@
 
 # CIFAR10 dataset consists of 50K training images. We define the batch size of 10 to load 5,000 batches of images.
 batch_size = 10
-number_of_labels = 10 
+number_of_labels = 10
 
-# Create an instance for training.  
-# When we run this code for the first time, the CIFAR10 train dataset will be downloaded locally. 
+# Create an instance for training.
+# When we run this code for the first time, the CIFAR10 train dataset will be downloaded locally.
 train_set =CIFAR10(root="./data",train=True,transform=transformations,download=True)
 
 # Create a loader for the training set which will read the data within batch size and put into memory.
 train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=0)
 print("The number of images in a training set is: ", len(train_loader)*batch_size)
 
 # Create an instance for testing, note that train is set to False.
-# When we run this code for the first time, the CIFAR10 test dataset will be downloaded locally. 
+# When we run this code for the first time, the CIFAR10 test dataset will be downloaded locally.
 test_set = CIFAR10(root="./data", train=False, transform=transformations, download=True)
 
-# Create a loader for the test set which will read the data within batch size and put into memory. 
+# Create a loader for the test set which will read the data within batch size and put into memory.
 # Note that both shuffle is set to false for the test loader.
 test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=0)
 print("The number of images in a test set is: ", len(test_loader)*batch_size)
@@ -49,12 +49,12 @@
 class Network(nn.Module):
     def __init__(self):
         super(Network, self).__init__()
-        
+
         self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=5, stride=1, padding=1)
         self.bn1 = nn.BatchNorm2d(12)
         self.conv2 = nn.Conv2d(in_channels=12, out_channels=12, kernel_size=5, stride=1, padding=1)
         self.bn2 = nn.BatchNorm2d(12)
-        
+
         self.pool = nn.MaxPool2d(2,2)
 
         self.conv4 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=5, stride=1, padding=1)
@@ -75,8 +75,9 @@ def forward(self, input):
 
         return output
 
-# Instantiate a neural network model 
+# Instantiate a neural network model
 model = Network()
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 # Define the loss function with Classification Cross-Entropy loss and an optimizer with Adam optimizer
 loss_fn = nn.CrossEntropyLoss()
@@ -89,34 +90,34 @@ def saveModel():
 
 # Function to test the model with the test dataset and print the accuracy for the test images
 def testAccuracy():
-    
+
     model.eval()
     accuracy = 0.0
     total = 0.0
-    
+
     with torch.no_grad():
-        
+
         for data in test_loader:
             images, labels = data
             # Run the model on the test set to predict labels
-            outputs = model(images)
+            outputs = model(images.to(device))
+            outputs = outputs.cpu()
             # The label with the highest energy will be our prediction
             _, predicted = torch.max(outputs.data, 1)
             total += labels.size(0)
             accuracy += (predicted == labels).sum().item()
-    
+
     # Compute the accuracy over all test images
     accuracy = (100 * accuracy / total)
-    return(accuracy)   
+    return(accuracy)
 
 
 # Training function. We simply have to loop over our data iterator, and feed the inputs to the network and optimize.
 def train(num_epochs):
-    
+
     best_accuracy = 0.0
 
     # Define your execution device
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     print("The model will be running on", device, "device")
     # Convert model parameters and buffers to CPU or Cuda
     model.to(device)
@@ -125,7 +126,7 @@ def train(num_epochs):
         running_loss = 0.0
 
         for i, (images, labels) in enumerate(train_loader, 0):
-            
+
             # wrap the inputs in Variable
             images = Variable(images.to(device))
             labels = Variable(labels.to(device))
@@ -143,8 +144,8 @@ def train(num_epochs):
 
             # Let's print statistics for every 1,000 images
             running_loss += loss.item()     # extract the loss value
-            if i % 1000 == 999:    
-                # print every 1000 (twice per epoch) 
+            if i % 1000 == 999:
+                # print every 1000 (twice per epoch)
                 print('[%d, %5d] loss: %.3f' %
                       (epoch + 1, i + 1, running_loss / 1000))
                 # zero the loss
@@ -153,7 +154,7 @@ def train(num_epochs):
         # Compute and print the average accuracy of this epoch when tested over all 10000 test images
         accuracy = testAccuracy()
         print('For epoch', epoch+1,'the test accuracy over the whole test set is %d %%' % (accuracy))
-        
+
         # We want to save the model if the accuracy is the best
         if accuracy > best_accuracy:
             saveModel()
@@ -170,37 +171,37 @@ def imageshow(img):
 
 # Function to test the model with a batch of images and show the labels predictions
 def testBatch():
-    # get batch of images from the test DataLoader  
+    # get batch of images from the test DataLoader
     images, labels = next(iter(test_loader))
 
     # show all images as one image grid
     imageshow(torchvision.utils.make_grid(images))
-   
-    # Show the real labels on the screen 
-    print('Real labels: ', ' '.join('%5s' % classes[labels[j]] 
+
+    # Show the real labels on the screen
+    print('Real labels: ', ' '.join('%5s' % classes[labels[j]]
                                for j in range(batch_size)))
-  
+
     # Let's see what if the model identifiers the  labels of those example
     outputs = model(images)
-    
+
     # We got the probability for every 10 labels. The highest (max) probability should be correct label
     _, predicted = torch.max(outputs, 1)
-    
+
     # Let's show the predicted labels on the screen to compare with the real ones
-    print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] 
+    print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]
                               for j in range(batch_size)))
- 
+
 
 #Function to Convert to ONNX
 def Convert_ONNX():
-    
+
     # set the model to inference mode
     model.eval()
-    
-    # Let's create a dummy input tensor 
+
+    # Let's create a dummy input tensor
     dummy_input = torch.randn(1, 3, 32, 32, requires_grad=True)
 
-    # Export the model  
+    # Export the model
     torch.onnx.export(model,                 # model being run
                   dummy_input,               # model input (or a tuple for multiple inputs)
                   "ImageClassifier.onnx",          # where to save the model (can be a file or file-like object)
@@ -234,17 +235,17 @@ def testClassess():
         print('Accuracy of %5s : %2d %%' % (
             classes[i], 100 * class_correct[i] / class_total[i]))
 
-    
+
 
 if __name__ == "__main__":
-    
+
     # Let's build our model
     train(5)
     print('Finished Training')
 
     # Test which classes performed well
     testAccuracy()
-    
+
     # Let's load the model we just created and test the accuracy per label
     model = Network()
     path = "myFirstModel.pth"
@@ -257,7 +258,3 @@ def testClassess():
 
     # Conversion to ONNX
     Convert_ONNX()
-        
-
-     
-