File size: 2,861 Bytes
deb7039 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import torch
import torchvision
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader
from torchvision import models
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import os
import os.path
import pickle
from typing import Any, Callable, Optional, Tuple
import torchvision.transforms.functional as TF
from tools.dataset import UpsideDownDataset
from train import train
from test import test
from observations import observations
import torchvision.models as models
import torch.optim as optim
import argparse
# setting up the device
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
def arg_parser():
parser = argparse.ArgumentParser(description="Upside Down Image Detector")
parser.add_argument('--epochs', type=int, default=10)
parser.add_argument('--pretrained', type=bool, default=True)
return parser.parse_args()
def main(args):
# transformations
train_transform = transforms.Compose([
transforms.ToTensor()
])
test_transform = transforms.Compose([
transforms.ToTensor()
])
# downloading CIFAR10 dataset
trainset = UpsideDownDataset(root='./data', train=True, download=True, transform=train_transform)
testset = UpsideDownDataset(root='./data', train=False, download=True, transform=test_transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, drop_last=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=128)
model = models.resnet18(pretrained=args.pretrained)
model.fc = torch.nn.Linear(512, 2)
optimizer = optim.SGD(model.parameters(), lr=0.01)
# criterion = torch.nn.BCEWithLogitsLoss()
criterion = torch.nn.CrossEntropyLoss()
model, criterion = model.to(DEVICE), criterion.to(DEVICE)
print("training started:")
loss_values = []
accuracy_values = []
for epoch in range(args.epochs):
print("Epoch {}".format(epoch+1))
epoch_loss = train(model, trainloader, optimizer, criterion, DEVICE)
loss_values.append(epoch_loss)
accuracy = test(model, testloader, criterion, DEVICE)
accuracy_values.append(accuracy)
plt.plot(np.array(loss_values), 'r')
plt.xlabel('epochs')
plt.ylabel('losses')
plt.savefig('./viz/loss_curve.png')
fig, ax = plt.subplots()
ax.plot(np.array(loss_values), color='red', label='losses')
ax.tick_params(axis='y', labelcolor='red')
ax2 = ax.twinx()
ax2.plot(np.array(accuracy_values), color='green', label='accuracy')
ax2.tick_params(axis='y', labelcolor='green')
# plt.xlabel('epochs')
plt.legend()
plt.savefig('./viz/final_curve.png')
print("accuracy: {}".format(accuracy))
observations(model, testloader)
if __name__ =='__main__':
args = arg_parser()
main(args)
|