image_classification/training.py at master · skbearman/image_classification · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
"""
training.py
Zhiang Chen, April 2020
"""

import torch
import torch.utils.data
import torchvision.datasets
import torch.nn as nn
import torchvision.transforms as transforms
from utils import *
import torchvision.models as models
from data import EurekaDataset

model_names = sorted(name for name in models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(models.__dict__[name]))

torch.manual_seed(0)

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])


eureka_normalize = transforms.Normalize(mean=[0.44, 0.50, 0.43],
                                     std=[0.26, 0.25, 0.26])

eureka_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    eureka_normalize,])


train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    normalize,])

test_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.ToTensor(),
    normalize,])

def neural_network(architecture, nm_classes, pretrained=True, change_last_layer=True):
    assert architecture in model_names
    print("=> creating model '{}'".format(architecture))
    model = models.__dict__[architecture](pretrained=pretrained)
    if change_last_layer:
        if architecture.startswith('densenet'):
            in_features = model.classifier.in_features
            model.classifier = nn.Linear(in_features=in_features, out_features=nm_classes)
        else:
            in_features = model.fc.in_features
            model.fc = nn.Linear(in_features=in_features, out_features=nm_classes)

    return model

def cifar10(root='./datasets/cifar10/', val=True):
    train = torchvision.datasets.CIFAR10(root, train=True, download=True, transform=train_transform)
    test = torchvision.datasets.CIFAR10(root, train=False, download=True, transform=test_transform)
    """
    if val:
        indices = torch.randperm(len(train)).tolist()
        train_set = torch.utils.data.Subset(train, indices[:-10000])
        val_set = torch.utils.data.Subset(train, indices[-10000:])
        return train_set, val_set, test
    """
    return train, test

def eureka():
	train = EurekaDataset('./datasets/Eureka/images0/','./datasets/Eureka/class.json', eureka_transform)
	test = EurekaDataset('./datasets/Eureka/images_test0/','./datasets/Eureka/class.json', eureka_transform)
	return train, test

if __name__ == '__main__':
    cuda = 'cuda:0'
    device = torch.device(cuda)
    nm_classes = 3
    train_dataset, test_dataset = eureka()

    train_dataloader = torch.utils.data.DataLoader(
        train_dataset, batch_size=4, shuffle=True, num_workers=8, collate_fn=collate_fn)
    test_dataloader = torch.utils.data.DataLoader(
        test_dataset, batch_size=4, shuffle=True, num_workers=8, collate_fn=collate_fn)

    model = neural_network('resnext101_32x8d', nm_classes)

    #if you want to load weight
	#model.load_state_dict(torch.load("trained_param_eureka_cls/epoch_0002.param"))
	#model.eval()

    model.to(device)

    criterion = nn.CrossEntropyLoss().to(device)
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=0.001, momentum=0.9, weight_decay=0.00001)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.65)

    init_epoch = 0
    num_epochs = 60
    print_freq = 100

    save_param = "trained_param3_resnext101/epoch_{:04d}.param".format(init_epoch)
    torch.save(model.state_dict(), save_param)

    for epoch in range(init_epoch, init_epoch + num_epochs):
        save_param = "trained_param3_resnext101/epoch_{:04d}.param".format(epoch)
        train(train_dataloader, model, criterion, optimizer, epoch, device, print_freq)
        lr_scheduler.step()
        #validate(test_dataloader, model, criterion, device)
        acc = test(model, test_dataset, device)
        print("acc: %f" % acc)
        torch.save(model.state_dict(), save_param)