ResNet 实现Cifar-10 识别以及一点思考

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

%matplotlib inline
plt.rcParams['figure.figsize'] = (10.0, 8.0) # set default size of plots
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'gray'

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),  #先四周填充0，再把图像随机裁剪成32*32
    transforms.RandomHorizontalFlip(),  #图像一半的概率翻转，一半的概率不翻转
    transforms.ToTensor(),
    transforms.Normalize((0, 0, 0), (1, 1, 1)), #R,G,B每层的归一化用到的均值和方差
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0, 0, 0), (1, 1, 1)),
])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,
                                          shuffle=True, num_workers=0)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=128,
                                         shuffle=False, num_workers=0)

1
2
3
4
5
6
7
8
9
class Flatten(nn.Module):
    def forward(self, x):
        N = x.shape[0]
        return x.view(N, -1)
class PrintParamShape(nn.Module):
    def forward(self, x):
        print(x.size())
        return x

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def DrawPlot(train_loss_summery, train_acc_summery, test_acc_summery):
    plt.subplot(2, 1, 1)
    plt.grid()
    plt.title('Training loss')
    plt.plot(train_loss_summery, 'o')
    plt.xlabel('Epoch')

    plt.subplot(2, 1, 2)
    plt.grid()
    plt.title('Accuracy')
    plt.plot(train_acc_summery, '-o', label='train')
    plt.plot(test_acc_summery, '-o', label='val')
    plt.plot([0.5] * len(test_acc_summery), 'k--')
    plt.xlabel('Epoch')
    plt.legend(loc='lower right')
    plt.gcf().set_size_inches(15, 12)
    plt.show()

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
def AdjustLearningRate(optimizer, epoch, iter_times):
    if iter_times == 32000 or iter_times == 48000 or iter_times == 52000 or iter_times == 56000:
        for param_group in optimizer.param_groups:
            param_group['lr'] *= 0.1
        print("change lr")

def TrainNet(net, datasetLoader, criterion, optimizer, device, epoch_total):
    train_loss_summery = []
    train_acc_summery = []
    test_acc_summery = []
    (trainloader, testloader) = datasetLoader
    running_loss = 0.0
    iter_times = 0
    for epoch in range(epoch_total + 1):
        train_loss_summery.append(running_loss)
        train_acc_summery.append(PredictNet(net, trainloader, device, 0))
        test_acc_summery.append(PredictNet(net, testloader, device, 0))
        print('epoch %d, loss %.3f, train_acc %.3f %%, test_acc %.3f %%' % (epoch, running_loss / 1e2, train_acc_summery[epoch],test_acc_summery[epoch]))
        if epoch == epoch_total or iter_times > 64000:
            break
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data[0].to(device), data[1].to(device)
            optimizer.zero_grad()

            output = net(inputs)        

            loss = criterion(output, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            iter_times += 1
            AdjustLearningRate(optimizer, epoch, iter_times)
        
    print('Finished Training')
    return (train_loss_summery, train_acc_summery, test_acc_summery)   

def PredictNet(net, testloader, device, num = 0, isPrint=False):
    total = 0
    correct = 0
    accuracy = 0.0
    with torch.no_grad():
        for data in testloader:
            inputs, labels = data[0].to(device), data[1].to(device)
            outputs = net(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            if total > num and num != 0:
                break
    accuracy = 100 * correct / total
    if isPrint == True:
        print('Accuracy of the network on the %d test images: %.3f %%' % (
            total, accuracy))
    return accuracy
         

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
def HeInitParam(layer):
#     return layer #如果不想初始化就把这个注释去掉
    if isinstance(layer, (torch.nn.Conv2d, torch.nn.Linear)):
        nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')
        nn.init.zeros_(layer.bias)
    return layer
class ResidualBlock(nn.Module):
    def __init__(self, params):
        super(ResidualBlock,self).__init__()
        self.ifPoolConnect = params['ifPoolConnect']
        self.input_channel = params['input_channel']
        self.repeat_arch = params['repeat_arch']#repeat_arch[i][j] means ith layer, j=0 means kernel size, else means feature map cnt
        self.repeat_times = params['repeat_times']
        
        partModule = None
        self.poolLayer = torch.nn.ModuleList()
        self.ResBlockModules = torch.nn.ModuleList()
        if self.ifPoolConnect == True: 
            self.poolLayer.append(HeInitParam(torch.nn.Conv2d(self.input_channel, self.repeat_arch[len(self.repeat_arch)-1][1], kernel_size=2, stride=2)))

        for i in range(self.repeat_times):
            partModule = torch.nn.ModuleList()
            if self.ifPoolConnect == True:    
                partModule.append(self.getConvBNReLu(self.repeat_arch[len(self.repeat_arch)-1][1], self.repeat_arch[0][1], self.repeat_arch[0][0], _padding=(self.repeat_arch[0][0]-1)//2))
            else:
                partModule.append(self.getConvBNReLu(self.input_channel, self.repeat_arch[0][1], self.repeat_arch[0][0], _padding=(self.repeat_arch[0][0]-1)//2))

            for partLayerCnter in range(1, len(self.repeat_arch)):
                if partLayerCnter == len(self.repeat_arch) - 1:
                    partModule.append(HeInitParam(torch.nn.Conv2d(self.repeat_arch[partLayerCnter-1][1], self.repeat_arch[partLayerCnter][1], self.repeat_arch[partLayerCnter][0], padding=(self.repeat_arch[partLayerCnter][0]-1)//2)))
                    partModule.append(torch.nn.BatchNorm2d(self.repeat_arch[partLayerCnter][1]))
                else:
                    partModule.append(self.getConvBNReLu(self.repeat_arch[partLayerCnter-1][1], self.repeat_arch[partLayerCnter][1], self.repeat_arch[partLayerCnter][0], (self.repeat_arch[partLayerCnter][0]-1)//2))
                
            self.ResBlockModules.append(partModule)
                                                     
    def getConvBNReLu(self, channel_in, channel_out, kernel_size, _padding=0):
        return torch.nn.Sequential(HeInitParam(torch.nn.Conv2d(channel_in, channel_out, kernel_size, padding=_padding)),
                                        torch.nn.BatchNorm2d(channel_out),
                                        torch.nn.ReLU(),)
    def forward(self, x):
        #This may not be right.
        if(self.ifPoolConnect == True):
            x = self.poolLayer[0](x)
        for i in range(self.repeat_times):
            _x = x.clone()
            for partLayerCnter in range(0, len(self.repeat_arch)):
                x = self.ResBlockModules[i][partLayerCnter](x)
            x = _x + x
            x = torch.nn.functional.relu(x)
                                                     
                                                     
        return x

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class ResidualNet(nn.Module):
    def __init__(self):
        super(ResidualNet,self).__init__()
        repeat_times = 7
        self.ResBlockTestParam1 = {'ifPoolConnect':False, 'input_channel':16, 
                        'repeat_arch':[[3,16], [3,16],],'repeat_times':repeat_times}
        self.ResBlockTestParam2 = {'ifPoolConnect':True, 'input_channel':16, 
                        'repeat_arch':[[3,32], [3,32],],'repeat_times':repeat_times}
        self.ResBlockTestParam3 = {'ifPoolConnect':True, 'input_channel':32, 
                        'repeat_arch':[[3,64], [3,64],],'repeat_times':repeat_times}
        self.ResNetModule = torch.nn.Sequential(HeInitParam(torch.nn.Conv2d(3, 16, kernel_size=3, padding = 1)),
                                                torch.nn.BatchNorm2d(16),
                                                torch.nn.ReLU(),
                                                ResidualBlock(self.ResBlockTestParam1),
                                                ResidualBlock(self.ResBlockTestParam2),
                                                ResidualBlock(self.ResBlockTestParam3),  
                                                torch.nn.AvgPool2d(8, 1),
                                                Flatten(),

                                                HeInitParam(torch.nn.Linear(64, 10)),)
    def forward(self,x):
        return self.ResNetModule(x)

1
2
3
4
5
6
7
8
9
criterion = nn.CrossEntropyLoss()
net = ResidualNet().to(device)

optimizer = optim.SGD(net.parameters(), lr=1e-1, momentum=0.9, weight_decay=5e-4)
(train_loss_summery, train_acc_summery, test_acc_summery) = TrainNet(net, (trainloader, testloader), criterion, optimizer, device, 200)

PATH = './cifar_ResNet_.pth'
torch.save(net.state_dict(), PATH)