File size: 3,090 Bytes
a189a79 |
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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms, datasets, models
from torchvision.utils import make_grid
import os
import time
from PIL import ImageFile
import math
from model import ConvolutionalNet
ImageFile.LOAD_TRUNCATED_IMAGES = True
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
train_transforms = transforms.Compose([
transforms.RandomRotation(10),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
train_dataset = datasets.ImageFolder(root='./data/train', transform=train_transforms)
test_dataset = datasets.ImageFolder(root='./data/test', transform=test_transform)
torch.manual_seed(42)
train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=10, shuffle=False)
class_names = train_dataset.classes
for images, labels in train_loader:
break
torch.manual_seed(101)
model = ConvolutionalNet()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
start_time = time.time()
epochs = 5
# BATCH LIMITS
max_trn_batch = 800
max_tst_batch = 300
train_losses = []
test_losses = []
train_correct = []
test_correct = []
for epoch in range(epochs):
trn_corr = 0
tst_corr = 0
for b, (X_train, y_train) in enumerate(train_loader):
# if b == max_trn_batch:
# break
y_pred = model(X_train)
loss = criterion(y_pred, y_train)
if b % 200 == 0:
print(f"Epoch: {epoch+1}/{epochs}\tBatch: {b+1}\tLoss: {loss.item()}")
predicted = torch.max(y_pred, 1)[1]
batch_corr = (predicted == y_train).sum()
trn_corr += batch_corr
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_losses.append(loss)
train_correct.append(trn_corr)
# TEST
with torch.no_grad():
for b, (X_test, y_test) in enumerate(test_loader):
# if b == max_tst_batch:
# break
try:
y_pred = model(X_test)
except:
print("Error testing images")
continue
loss = criterion(y_pred, y_test)
predicted = torch.max(y_pred, 1)[1]
batch_corr = (predicted == y_test).sum()
tst_corr += batch_corr
test_losses.append(loss)
test_correct.append(tst_corr)
end_time = time.time()
total_time = end_time - start_time
print(f"Time taken: minutes: {math.floor(total_time / 60)} seconds: {math.floor(total_time % 60)}")
torch.save(model.state_dict(), 'model.pt')
plt.plot([x.detach().numpy() for x in train_losses], label='train loss')
plt.plot(test_losses, label='test loss')
plt.legend()
plt.plot()
plt.plot([t/80 for t in train_correct], label='train accuracy')
plt.plot([t/30 for t in test_correct], label='test accuracy')
plt.legend()
plt.plot()
print(f'Accuracy: {100*test_correct[-1].item()/1000}')
|