-
Notifications
You must be signed in to change notification settings - Fork 22
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #300 from bformby/main
Updating fabric_gpu example notebook
- Loading branch information
Showing
2 changed files
with
182 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
108 changes: 108 additions & 0 deletions
108
fabric_examples/fablib_api/fabric_all_gpus/pytorch_example.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,108 @@ | ||
| import torch | ||
| import torchvision | ||
| import torchvision.transforms as transforms | ||
| import torch.nn as nn | ||
| import torch.nn.functional as F | ||
| import torch.optim as optim | ||
|
|
||
| device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') | ||
| # Assuming that we are on a CUDA machine, this should print a CUDA device: | ||
| print("Device: ",device) | ||
|
|
||
| # --- Load and normalize CIFAR10 --- | ||
| transform = transforms.Compose( | ||
| [transforms.ToTensor(), | ||
| transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) | ||
|
|
||
| batch_size = 4 | ||
|
|
||
| trainset = torchvision.datasets.CIFAR10(root='./data', train=True, | ||
| download=True, transform=transform) | ||
| trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, | ||
| shuffle=True, num_workers=2) | ||
|
|
||
| testset = torchvision.datasets.CIFAR10(root='./data', train=False, | ||
| download=True, transform=transform) | ||
| testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, | ||
| shuffle=False, num_workers=2) | ||
|
|
||
| classes = ('plane', 'car', 'bird', 'cat', | ||
| 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') | ||
|
|
||
| # --- Define a Convolutional Neural Network --- | ||
| class Net(nn.Module): | ||
| def __init__(self): | ||
| super().__init__() | ||
| self.conv1 = nn.Conv2d(3, 6, 5) | ||
| self.pool = nn.MaxPool2d(2, 2) | ||
| self.conv2 = nn.Conv2d(6, 16, 5) | ||
| self.fc1 = nn.Linear(16 * 5 * 5, 120) | ||
| self.fc2 = nn.Linear(120, 84) | ||
| self.fc3 = nn.Linear(84, 10) | ||
|
|
||
| def forward(self, x): | ||
| x = self.pool(F.relu(self.conv1(x))) | ||
| x = self.pool(F.relu(self.conv2(x))) | ||
| x = torch.flatten(x, 1) # flatten all dimensions except batch | ||
| x = F.relu(self.fc1(x)) | ||
| x = F.relu(self.fc2(x)) | ||
| x = self.fc3(x) | ||
| return x | ||
|
|
||
|
|
||
| net = Net().to(device) | ||
|
|
||
| # --- Define a Loss function and optimizer --- | ||
| criterion = nn.CrossEntropyLoss() | ||
| optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) | ||
|
|
||
| # --- Train the network --- | ||
| for epoch in range(2): # loop over the dataset multiple times | ||
|
|
||
| running_loss = 0.0 | ||
| for i, data in enumerate(trainloader, 0): | ||
| # get the inputs; data is a list of [inputs, labels] | ||
| inputs, labels = data[0].to(device), data[1].to(device) | ||
|
|
||
| # zero the parameter gradients | ||
| optimizer.zero_grad() | ||
|
|
||
| # forward + backward + optimize | ||
| outputs = net(inputs) | ||
| loss = criterion(outputs, labels) | ||
| loss.backward() | ||
| optimizer.step() | ||
|
|
||
| # print statistics | ||
| running_loss += loss.item() | ||
| if i % 2000 == 1999: # print every 2000 mini-batches | ||
| print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}') | ||
| running_loss = 0.0 | ||
|
|
||
| print('Finished Training') | ||
| PATH = './cifar_net.pth' | ||
| torch.save(net.state_dict(), PATH) | ||
|
|
||
| # --- Test the network on the test data --- | ||
| dataiter = iter(testloader) | ||
| images, labels = next(dataiter) | ||
|
|
||
| net = Net() | ||
| net.load_state_dict(torch.load(PATH)) | ||
| outputs = net(images) | ||
| _, predicted = torch.max(outputs, 1) | ||
|
|
||
| correct = 0 | ||
| total = 0 | ||
| # since we're not training, we don't need to calculate the gradients for our outputs | ||
| with torch.no_grad(): | ||
| for data in testloader: | ||
| images, labels = data | ||
| # calculate outputs by running images through the network | ||
| outputs = net(images) | ||
| # the class with the highest energy is what we choose as prediction | ||
| _, predicted = torch.max(outputs.data, 1) | ||
| total += labels.size(0) | ||
| correct += (predicted == labels).sum().item() | ||
|
|
||
| print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %') |