157 lines
5.1 KiB
Python
157 lines
5.1 KiB
Python
import torch
|
|
import torchvision
|
|
import torchvision.transforms as transforms
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
import torch.optim as optim
|
|
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
#import dla
|
|
import vgg
|
|
|
|
|
|
|
|
EPOCHS = 40
|
|
|
|
class CnnBlock(nn.Module):
|
|
def __init__(self, in_channels=3, out_channels=16):
|
|
super(CnnBlock, self).__init__()
|
|
self.conv1 = nn.Conv2d(in_channels, 2*in_channels, 3, 1)
|
|
self.conv2 = nn.Conv2d(2*in_channels, out_channels, 3, 1)
|
|
self.bn1 = nn.BatchNorm2d(2*in_channels)
|
|
self.bn2 = nn.BatchNorm2d(out_channels)
|
|
|
|
def forward(self, x):
|
|
x = self.conv1(x)
|
|
x = self.bn1(x)
|
|
x = F.relu(x)
|
|
x = self.conv2(x)
|
|
x = self.bn2(x)
|
|
x = F.relu(x)
|
|
return x
|
|
|
|
class CifarCNN(nn.Module):
|
|
def __init__(self):
|
|
super(CifarCNN, self).__init__()
|
|
self.block1 = CnnBlock(3, 16)
|
|
self.block2 = CnnBlock(16, 64)
|
|
self.block3 = CnnBlock(64, 128)
|
|
self.dropout1 = nn.Dropout(0.25)
|
|
self.dropout2 = nn.Dropout(0.5)
|
|
self.fc1 = nn.Linear(128, 10)
|
|
|
|
def forward(self, x):
|
|
x = self.block1(x)
|
|
x = self.dropout1(x)
|
|
x = F.max_pool2d(x,2)
|
|
x = self.block2(x)
|
|
x = F.max_pool2d(x,2)
|
|
x = self.dropout2(x)
|
|
x = self.block3(x)
|
|
x = torch.flatten(x,1)
|
|
x = self.fc1(x)
|
|
output = F.log_softmax(x, dim=1)
|
|
return output
|
|
|
|
def train(model, trainloader, device, optimizer, criterion, epoch):
|
|
running_loss = 0.0
|
|
for i, [data, target] in enumerate(trainloader, 0):
|
|
data, target = data.to(device), target.to(device)
|
|
# zero the parameter gradients
|
|
optimizer.zero_grad()
|
|
|
|
# forward + backward + optimize
|
|
output = model(data)
|
|
loss = criterion(output, target)
|
|
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
|
|
|
|
|
|
def test(model, testloader, device, classes):
|
|
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, target in testloader:
|
|
data, target = data.to(device), target.to(device)
|
|
# calculate outputs by running images through the network
|
|
output = model(data)
|
|
# the class with the highest energy is what we choose as prediction
|
|
_, predicted = torch.max(output.data, 1)
|
|
total += target.size(0)
|
|
correct += (predicted == target).sum().item()
|
|
|
|
print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')
|
|
|
|
|
|
# prepare to count predictions for each class
|
|
correct_pred = {classname: 0 for classname in classes}
|
|
total_pred = {classname: 0 for classname in classes}
|
|
|
|
# again no gradients needed
|
|
with torch.no_grad():
|
|
for data, target in testloader:
|
|
data, target = data.to(device), target.to(device)
|
|
output = model(data)
|
|
_, predictions = torch.max(output, 1)
|
|
# collect the correct predictions for each class
|
|
for label, prediction in zip(target, predictions):
|
|
if label == prediction:
|
|
correct_pred[classes[label]] += 1
|
|
total_pred[classes[label]] += 1
|
|
|
|
|
|
# print accuracy for each class
|
|
for classname, correct_count in correct_pred.items():
|
|
accuracy = 100 * float(correct_count) / total_pred[classname]
|
|
print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')
|
|
|
|
|
|
def main():
|
|
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')
|
|
|
|
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
|
model = vgg.VGG('VGG16').to(device)
|
|
|
|
optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=0.9)
|
|
criterion = nn.CrossEntropyLoss()
|
|
|
|
for epoch in range(EPOCHS):
|
|
train(model, trainloader, device, optimizer, criterion, epoch)
|
|
test(model, testloader, device, classes)
|
|
|
|
PATH = './cifar_vgg.pth'
|
|
torch.save(model.state_dict(), PATH)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|