main.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd.variable import Variable
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets.cifar import CIFAR10
from model import ResNet
from tqdm import tqdm


def main():
    batch_size = 32
    model = ResNet(pretrained=True, num_class=10)
    print(model.parameters())
    weight_decay = 5e-4
    optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=weight_decay)
    criterion = nn.NLLLoss2d()

    transform = transforms.Compose([transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    train_data = CIFAR10(root='../datasets/', train=True, transform=transform, download=True)
    test_data = CIFAR10(root='../datasets/', train=False, transform=transform, download=True)

    train_loader = DataLoader(train_data, batch_size=32, shuffle=True, num_workers=4)
    val_loader = DataLoader(test_data, batch_size=32, shuffle=False, num_workers=4)

    pm = ProcessManage(model, criterion, optimizer)
    pm.run(train_loader, val_loader)


class ProcessManage(object):
    def __init__(self, model, criterion, optimizer):
        self.model = model
        self.criterion = criterion
        self.min_val_loss=100
        self.optimizer = optimizer
    
    def run(self, train_loader, val_loader):
        for e in tqdm(range(1, 5 + 1)):
            self.train(train_loader)
            self.validation(val_loader, e)

    def train(self, train_loader):
        self.model.train()
        for i, data in enumerate(train_loader):
            inputs, label = data
            inputs, label = Variable(inputs).cuda(), Variable(label).cuda()
            output = self.model(inputs)
            loss = self.criterion(output, label)
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
    
    def validation(self, val_loader, epoch):
        model.eval()
        val_loss = 0
        correct = 0
        is_best = False
    
        for i, data in enumerate(val_loader):
            inputs, label = data
            inputs, label = Variable(inputs, volatile=True).cuda(), Variable(label).cuda()
            output = self.model(data)
            val_loss += self.criterion(output, target, size_average=False).data[0]
            pred = output.data.max(1, keepdim=True)[1] 
            correct += pred.eq(target.data.view_as(pred)).long().cpu().sum()
    
        val_loss /= len(val_loader.dataset)

        if val_loss < min_val_loss:
            is_best = True
            save_checkpoint(
                {'epoch', epoch,
                 'state', self.model.state_dict
                }
            , is_best)
            self.min_val_loss = val_loss

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


def save_checkpoint(state, is_best, filename='results/checkpoint.pth.tar'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'results/model_best.pth.tar')


if __name__ == '__main__':
    main()

model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models.resnet import resnet50

class ResNet(nn.Module):
    def __init__(self, pretrained=False, num_class=10):
        super().__init__()
        resnet = resnet50(pretrained=pretrained)
        self.layers = nn.Sequential(*list(resnet.children())[:-2])
        self.pool = nn.AvgPool2d(kernel_size=1)
        self.fc = nn.Linear(in_features=2048, out_features=num_class)
        if pretrained is not True:
            for param in resnet.parameters():
                param.requires_grad = False
    
    def forward(self, inputs, target=None):
        h = self.layers(inputs)
        h = self.pool(h)
        out = h.view(h.size(0), -1)
        out = self.fc(out)
        return out