From b5157f844172966ff7688ad0860765a42f7a217f Mon Sep 17 00:00:00 2001 From: Chris Lovett Date: Thu, 17 Nov 2022 09:23:33 -0800 Subject: [PATCH] fix sample so it works on the GPU (#505) --- .../PyTorchTraining.py | 77 +++++++++---------- 1 file changed, 37 insertions(+), 40 deletions(-) diff --git a/Samples/Tutorial Samples/PyTorch Image Classification/PyTorchTraining - Image Classification/PyTorchTraining.py b/Samples/Tutorial Samples/PyTorch Image Classification/PyTorchTraining - Image Classification/PyTorchTraining.py index d88f0c5b..cad462ee 100644 --- a/Samples/Tutorial Samples/PyTorch Image Classification/PyTorchTraining - Image Classification/PyTorchTraining.py +++ b/Samples/Tutorial Samples/PyTorch Image Classification/PyTorchTraining - Image Classification/PyTorchTraining.py @@ -21,10 +21,10 @@ transformations = transforms.Compose([ # 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. @@ -32,10 +32,10 @@ train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_wo 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 @@ classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship' 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 @@ class Network(nn.Module): 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 @@ if __name__ == "__main__": # Conversion to ONNX Convert_ONNX() - - - - \ No newline at end of file