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) | |