fix sample so it works on the GPU (#505)
This commit is contained in:
Chris Lovett
GitHub
Родитель
673410d94f
Коммит
b5157f8441
|
|
@ -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()
|
||||
|
||||
|
||||
|
||||
|
||||
Загрузка…
Ссылка в новой задаче