|
import os
|
|
import time
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torchvision.datasets as datasets
|
|
import torchvision.transforms as transforms
|
|
from torch.utils.data import DataLoader
|
|
import torch.optim as optim
|
|
|
|
from tqdm import tqdm
|
|
|
|
|
|
def get_mean_std(loader):
|
|
'''
|
|
Calculates mean and std of input images.
|
|
|
|
Args:
|
|
loader (torch.DataLoader): Loader with images
|
|
|
|
Returns:
|
|
mean (torch.Tensor): Mean of images in loader
|
|
std (torch.Tensor): Standard deviation of images in loader
|
|
'''
|
|
channels_sum, channels_squared_sum, num_batches = 0, 0, 0
|
|
|
|
for data, _ in loader:
|
|
channels_sum += torch.mean(data, dim=[0,2,3])
|
|
channels_squared_sum += torch.mean(data**2, dim=[0,2,3])
|
|
num_batches += 1
|
|
|
|
mean = channels_sum/num_batches
|
|
std = (channels_squared_sum/(num_batches-mean**2))**0.5
|
|
|
|
return mean, std
|
|
|
|
|
|
class Net(nn.Module):
|
|
'''
|
|
model definition
|
|
'''
|
|
|
|
def __init__(self):
|
|
super(Net, self).__init__()
|
|
|
|
self.layer1 = nn.Sequential(
|
|
nn.Conv2d(1, 32, kernel_size=5),
|
|
nn.ReLU(),
|
|
)
|
|
self.layer2 = nn.Sequential(
|
|
nn.Conv2d(32, 32, kernel_size=5, bias=False),
|
|
nn.BatchNorm2d(32),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d((2, 2)),
|
|
nn.Dropout2d(0.25),
|
|
)
|
|
self.layer3 = nn.Sequential(
|
|
nn.Conv2d(32, 64, kernel_size=3),
|
|
nn.ReLU(),
|
|
)
|
|
self.layer4 = nn.Sequential(
|
|
nn.Conv2d(64, 64, kernel_size=3, bias=False),
|
|
nn.BatchNorm2d(64),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d((2, 2)),
|
|
nn.Dropout2d(0.25),
|
|
nn.Flatten(),
|
|
)
|
|
self.layer5 = nn.Sequential(
|
|
nn.Linear(576, 256, bias=False),
|
|
nn.BatchNorm1d(256),
|
|
nn.ReLU(),
|
|
)
|
|
self.layer6 = nn.Sequential(
|
|
nn.Linear(256, 128, bias=False),
|
|
nn.BatchNorm1d(128),
|
|
nn.ReLU(),
|
|
)
|
|
self.layer7 = nn.Sequential(
|
|
nn.Linear(128, 84, bias=False),
|
|
nn.BatchNorm1d(84),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.25),
|
|
)
|
|
self.layer8 = nn.Sequential(
|
|
nn.Linear(84, 10),
|
|
nn.LogSoftmax(dim=1),
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = transforms.Normalize(mean, std)(x)
|
|
x = self.layer1(x)
|
|
x = self.layer2(x)
|
|
x = self.layer3(x)
|
|
x = self.layer4(x)
|
|
x = self.layer5(x)
|
|
x = self.layer6(x)
|
|
x = self.layer7(x)
|
|
x = self.layer8(x)
|
|
|
|
return x
|
|
|
|
|
|
def download_data():
|
|
|
|
transform = transforms.Compose([transforms.ToTensor(), transforms.RandomAffine(degrees=10, translate=(0.1,0.1))])
|
|
train_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
|
|
train_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=True, pin_memory=True)
|
|
|
|
|
|
val_data = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
|
|
val_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=False, pin_memory=True)
|
|
|
|
|
|
mean, std = get_mean_std(train_loader)
|
|
|
|
|
|
|
|
if torch.cuda.is_available():
|
|
dev = "cuda:0"
|
|
else:
|
|
dev = "cpu"
|
|
|
|
device = torch.device(dev)
|
|
|
|
|
|
def run_model():
|
|
|
|
model = Net().to(device=device)
|
|
optimizer = optim.Adam(model.parameters(), lr=0.1)
|
|
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=2)
|
|
criterion = nn.NLLLoss()
|
|
|
|
|
|
for epoch in range(30):
|
|
print(f"\nEpoch {epoch+1}/{30}.")
|
|
|
|
|
|
model.train()
|
|
|
|
total_train_loss = 0
|
|
total_correct = 0
|
|
|
|
for i, (images, labels) in enumerate(tqdm(train_loader)):
|
|
images = images.to(device)
|
|
labels = labels.to(device)
|
|
|
|
optimizer.zero_grad()
|
|
|
|
outputs = model(images)
|
|
|
|
loss = criterion(outputs, labels)
|
|
total_train_loss += loss.item()
|
|
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
|
|
outputs_probs = nn.functional.softmax(
|
|
outputs, dim=1)
|
|
for idx, preds in enumerate(outputs_probs):
|
|
|
|
if labels[idx] == torch.argmax(preds.data):
|
|
total_correct += 1
|
|
|
|
train_loss = total_train_loss/(i+1)
|
|
train_accuracy = total_correct/len(train_data)
|
|
|
|
print(f"Train set:- Loss: {train_loss}, Accuracy: {train_accuracy}.")
|
|
|
|
|
|
if not os.path.exists("./saved_models"):
|
|
os.mkdir("./saved_models")
|
|
torch.save(model.state_dict(), f"./saved_models/mnist-cnn-{time.time()}.pt")
|
|
|
|
|
|
model.eval()
|
|
|
|
total_val_loss = 0
|
|
total_correct = 0
|
|
|
|
with torch.no_grad():
|
|
for i, (images, labels) in enumerate(tqdm(val_loader)):
|
|
images = images.to(device)
|
|
labels = labels.to(device)
|
|
|
|
outputs = model(images)
|
|
|
|
loss = criterion(outputs, labels)
|
|
total_val_loss += loss.item()
|
|
|
|
outputs_probs = nn.functional.softmax(outputs, dim=1)
|
|
for idx, preds in enumerate(outputs_probs):
|
|
if labels[idx] == torch.argmax(preds.data):
|
|
total_correct += 1
|
|
|
|
val_loss = total_val_loss/(i+1)
|
|
val_accuracy = total_correct/len(val_data)
|
|
|
|
print(f"Val set:- Loss: {val_loss}, Accuracy: {val_accuracy}.")
|
|
|
|
|
|
scheduler.step(val_loss)
|
|
|