Spaces:
Runtime error
Runtime error
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]) # mean across [no. of examples, height, width] | |
channels_squared_sum += torch.mean(data**2, dim=[0,2,3]) # squared mean across [no. of examples, height, width] | |
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 | |
# downloads and loads MNIST train set | |
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) | |
# downloads and loads MNIST test set | |
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) | |
# uses GPU if available | |
if torch.cuda.is_available(): | |
dev = "cuda:0" | |
else: | |
dev = "cpu" | |
device = torch.device(dev) | |
# gets mean and std of dataset | |
mean, std = get_mean_std(train_loader) | |
def run_model(): | |
# defines parameters | |
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() | |
# iterates through epochs | |
for epoch in range(30): | |
print(f"\nEpoch {epoch+1}/{30}.") | |
# train loop | |
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() | |
# Calculates train accuracy | |
outputs_probs = nn.functional.softmax( | |
outputs, dim=1) # gets probabilities | |
for idx, preds in enumerate(outputs_probs): | |
# if label with max probability matches true label | |
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}.") | |
# saves model state | |
if not os.path.exists("./saved_models"): | |
os.mkdir("./saved_models") | |
torch.save(model.state_dict(), f"./saved_models/mnist-cnn-{time.time()}.pt") | |
# val loop | |
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}.") | |
# adjusts lr | |
scheduler.step(val_loss) | |