|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
import torch.optim as optim |
|
import torchvision.transforms as transforms |
|
from torch.utils.data import DataLoader, Dataset |
|
|
|
from cifar10 import Cifar10 |
|
from CustomCIFAR10Dataset import CustomCIFAR10Dataset |
|
|
|
|
|
|
|
cifar10_builder = Cifar10() |
|
|
|
cifar10_builder.download_and_prepare() |
|
|
|
|
|
train_data = cifar10_builder.as_dataset(split='train') |
|
test_data = cifar10_builder.as_dataset(split='test') |
|
|
|
train_images = train_data["img"] |
|
train_labels = train_data["label"] |
|
|
|
test_images = test_data["img"] |
|
test_labels = test_data["label"] |
|
|
|
|
|
classes = ("airplane", "automobile", "bird", "cat", "deer", |
|
"dog", "frog", "horse", "ship", "truck") |
|
|
|
|
|
|
|
batch_size = 128 |
|
|
|
|
|
img_size = 32 |
|
|
|
|
|
nc = 3 |
|
|
|
|
|
output = len(classes) |
|
|
|
|
|
ngpu = 1 |
|
|
|
|
|
nw = 0 |
|
|
|
|
|
num_epochs = 33 |
|
|
|
|
|
learning_rate = 0.0022 |
|
|
|
|
|
device = torch.device("cuda:0" if (torch.cuda.is_available()) and (ngpu > 0) else "cpu") |
|
|
|
|
|
transform = transforms.Compose([ |
|
transforms.Resize((32, 32)), |
|
transforms.RandomResizedCrop(32, scale=(0.8, 0.8)), |
|
transforms.RandomHorizontalFlip(), |
|
transforms.ToTensor(), |
|
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), |
|
]) |
|
|
|
|
|
train_ds = CustomCIFAR10Dataset(train_data["img"], train_data["label"], transform=transform) |
|
test_ds = CustomCIFAR10Dataset(test_data["img"], test_data["label"], transform=transform) |
|
|
|
|
|
train_loader = DataLoader(train_ds, batch_size, shuffle=True, num_workers=nw) |
|
test_loader = DataLoader(test_ds, batch_size, shuffle=True, num_workers=nw) |
|
|
|
|
|
|
|
class Net(nn.Module): |
|
def __init__(self): |
|
super(Net, self).__init__() |
|
|
|
self.network = nn.Sequential( |
|
|
|
nn.Conv2d(nc, 16, kernel_size=3, stride=1, padding=1), |
|
|
|
nn.BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), |
|
nn.ReLU(), |
|
nn.MaxPool2d(kernel_size=1, stride=1), |
|
nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1), |
|
nn.BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), |
|
nn.ReLU(), |
|
nn.MaxPool2d(1, 1), |
|
nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1), |
|
nn.BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), |
|
nn.ReLU(), |
|
nn.MaxPool2d(kernel_size=1, stride=1), |
|
nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1), |
|
nn.BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), |
|
nn.ReLU(), |
|
nn.MaxPool2d(kernel_size=1, stride=1), |
|
nn.Conv2d(32, 16, kernel_size=3, stride=1, padding=1), |
|
|
|
|
|
nn.MaxPool2d(kernel_size=2, stride=2), |
|
|
|
nn.Flatten(), |
|
nn.Linear(16 * 16 * 16, 64), |
|
nn.Linear(64, 128), |
|
nn.Dropout(0.2), |
|
nn.Linear(128, 16), |
|
nn.Linear(16, output), |
|
) |
|
|
|
|
|
def forward(self, x): |
|
return self.network(x) |
|
|
|
|
|
model = Net() |
|
|
|
print("model parameters: ", model.parameters) |
|
|
|
|
|
criterion = nn.CrossEntropyLoss() |
|
|
|
optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=0.001) |
|
|
|
|
|
model.to(device) |
|
|
|
|
|
accuracyArr = [] |
|
|
|
print("started training") |
|
for epoch in range(num_epochs): |
|
model.train() |
|
|
|
running_loss = 0.0 |
|
for batch_idx, (images, labels) in enumerate(train_loader): |
|
inputs = images.to(device) |
|
labels = labels.to(device) |
|
|
|
|
|
optimizer.zero_grad() |
|
|
|
|
|
predictions = model(inputs) |
|
|
|
|
|
loss = criterion(predictions, labels) |
|
|
|
|
|
loss.backward() |
|
|
|
|
|
optimizer.step() |
|
|
|
|
|
running_loss += loss.item() |
|
|
|
print(f"epoch: {epoch + 1}/{num_epochs} Loss: {running_loss}") |
|
|
|
|
|
model.eval() |
|
correct = 0 |
|
total = 0 |
|
|
|
with torch.no_grad(): |
|
for batch_idx, (images, labels) in enumerate(test_loader): |
|
images = images.to(device) |
|
labels = labels.to(device) |
|
|
|
|
|
predictions = model(images) |
|
|
|
|
|
|
|
_, predicted = torch.max(predictions.data, 1) |
|
|
|
|
|
total += labels.size(0) |
|
|
|
|
|
correct += (predicted == labels).sum().item() |
|
|
|
|
|
accuracy = correct/total |
|
accuracyArr.append(accuracy) |
|
print(f"Accuracy on the test dataset: {accuracy:.2%}, epoch: {epoch + 1}") |
|
|
|
|
|
|
|
model.eval() |
|
correct = 0 |
|
total = 0 |
|
|
|
with torch.no_grad(): |
|
for batch_idx, (images, labels) in enumerate(test_loader): |
|
images = images.to(device) |
|
labels = labels.to(device) |
|
|
|
|
|
predictions = model(images) |
|
|
|
|
|
|
|
_, predicted = torch.max(predictions.data, 1) |
|
|
|
|
|
total += labels.size(0) |
|
|
|
|
|
correct += (predicted == labels).sum().item() |
|
|
|
|
|
accuracy = correct/total |
|
print(f"Accuracy on the test dataset: {accuracy:.2%}") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|