p1_train.py

# -*- coding: utf-8 -*-
"""P1.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1VVUSPsRQphAM41joL4gkl96anZHYiTIa
"""

!gdown --id '1CIvfO8rDMq5-3vmi0c6mrhDfzqrAZDem' --output hw1_data.zip
! unzip hw1_data.zip

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import torchvision.models as models
from torch.utils.data import Dataset, DataLoader
import glob
import os
import numpy as np
from PIL import Image
import torchvision.models as models
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE

"""Data set"""

class VGG(Dataset):
  def __init__(self, root, transform=None):
    """ Intialize the VGG dataset """
    self.images = None
    self.labels = None
    self.filenames = []
    self.root = root
    self.transform = transform

    # read filenames
    filenames = glob.glob(os.path.join(root, '*.png'))
    for fn in filenames:
      file = fn.split('/')
      file = file[-1]
      label = file.split('_')
      label = int(label[0])
      self.filenames.append((fn, label)) # (file, label) pair
                
    self.len = len(self.filenames)
                              
  def __getitem__(self, index):
    """ Get a sample from the dataset """
    image_fn, label = self.filenames[index]
    image = Image.open(image_fn)
            
    if self.transform is not None:
      image = self.transform(image)

    return image, label

  def __len__(self):
    """ Total number of samples in the dataset """
    return self.len

transformdata = {
  'train': transforms.Compose([
  transforms.Resize((224,224)),
  transforms.RandomHorizontalFlip(),
  transforms.ToTensor(),
  transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
  'val': transforms.Compose([
  transforms.Resize((224,224)),
  transforms.ToTensor(),
  transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
}

trainset = VGG(root='p1_data/train_50', transform=transformdata['train'])
testset = VGG(root='p1_data/val_50', transform=transformdata['val'])
trainset_loader = DataLoader(trainset, batch_size=50, shuffle=True, num_workers=1)
testset_loader = DataLoader(testset, batch_size=50, shuffle=False, num_workers=1)

print('# images in trainset:', len(trainset))
print('# images in testset:', len(testset))
dataiter = iter(trainset_loader)
images, labels = dataiter.next()

print('Image tensor in each batch:', images.shape, images.dtype)
print('Label tensor in each batch:', labels.shape, labels.dtype)

def imshow(img):
  npimg = img.numpy()
  plt.imshow(np.transpose(npimg, (1, 2, 0)))
    
# show images
imshow(torchvision.utils.make_grid(images))
# print labels
print('Labels:')
print(' '.join('%5s' % labels[j] for j in range(10)))

# Use GPU if available, otherwise stick with cpu
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
print('Device used:', device)

VGG16_model = models.vgg16_bn(pretrained = True)
VGG16_model.classifier._modules['6'] = nn.Linear(4096, 50)
VGG16_model.to(device)

def save_checkpoint(checkpoint_path, model, optimizer):
  state = {'state_dict': model.state_dict(),
        'optimizer' : optimizer.state_dict()}
  torch.save(state, checkpoint_path)
  print('model saved to %s' % checkpoint_path)
    
def load_checkpoint(checkpoint_path, model, optimizer):
  state = torch.load(checkpoint_path)
  model.load_state_dict(state['state_dict'])
  optimizer.load_state_dict(state['optimizer'])
  print('model loaded from %s' % checkpoint_path)

def train(model, epoch, save_interval, log_interval=100):
  optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
  criterion = nn.CrossEntropyLoss()
  model.train()  # Important: set training mode
    
  iteration = 0
  for ep in range(epoch):
    for batch_idx, (data, target) in enumerate(trainset_loader):
      data, target = data.to(device), target.to(device)
      optimizer.zero_grad()
      output = model(data)
      loss = criterion(output, target)
      loss.backward()
      optimizer.step()
            
      if iteration % log_interval == 0:
        print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
            ep, batch_idx * len(data), len(trainset_loader.dataset),
            100. * batch_idx / len(trainset_loader), loss.item()))
      if iteration % save_interval == 0 and iteration > 0:
        save_checkpoint('VGG16-%i.pth' % iteration, model, optimizer)
        test(model)
      iteration += 1
            
    test(model) # Evaluate at the end of each epoch
  save_checkpoint('VGG16-%i.pth' % iteration, model, optimizer)

def test(model):
  criterion = nn.CrossEntropyLoss()
  model.eval()  # Important: set evaluation mode
  test_loss = 0
  correct = 0
  with torch.no_grad(): # This will free the GPU memory used for back-prop
    for data, target in testset_loader:
      data, target = data.to(device), target.to(device)
      output = model(data)
      test_loss += criterion(output, target).item() # sum up batch loss
      pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability
      correct += pred.eq(target.view_as(pred)).sum().item()

  test_loss /= len(testset_loader.dataset)
  print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
      test_loss, correct, len(testset_loader.dataset),
      100. * correct / len(testset_loader.dataset)))

train(VGG16_model, 20, 2000)

testmodel = models.vgg16_bn(pretrained = True)
testmodel.classifier._modules['6'] = nn.Linear(4096, 50)
testmodel.to(device)
state = torch.load('VGG16-9000.pth')
testmodel.load_state_dict(state['state_dict'])
torch.save(testmodel.state_dict(), 'mymodel.pth')