【PyTorch】：ResNet34实现

# Pytorch 0.4.0 ResNet34实现cifar10分类.
# @Time: 2018/6/17
# @Author: xfLi

import torchvision as tv
import torch as t
import torchvision.transforms as transforms
from torch import nn
from torch.utils.data import DataLoader
import torch.nn.functional as F
t.set_num_threads(8)


class ResidualBloak(nn.Module):
    #残差块
    def __init__(self, inchannel, outchannel, stride=1, shortcut=None):
        super(ResidualBloak, self).__init__()
        self.left = nn.Sequential(
            nn.Conv2d(inchannel, outchannel, 3, stride, 1, bias=False),
            nn.BatchNorm2d(outchannel),
            nn.ReLU(inplace=True),
            nn.Conv2d(outchannel, outchannel, 3, 1, 1, bias=False),
            nn.BatchNorm2d(outchannel))
        self.right = shortcut

    def forward(self, x):
        out = self.left(x)
        residual = x if self.right is None else self.right(x)
        out += residual
        return F.relu(out)

class ResNet34(nn.Module):
    #  实现主module：ResNet34  
    #  ResNet34 包含多个layer，每个layer又包含多个residual block  
    #  用子module来实现residual block，用_make_layer函数来实现layer 
    def __init__(self, num_classes):
        super(ResNet34, self).__init__()
        #前几层图像转换
        self.pre = nn.Sequential(
            nn.Conv2d(3, 16, 3, 1, 1, bias=False),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2, 1))

        # 重复的layer，分别有3，4，6，3个residual block
        self.layer1 = self._make_layer(16, 16, 3, stride=1)
        self.layer2 = self._make_layer(16, 32, 4, stride=1)
        self.layer3 = self._make_layer(32, 64, 6, stride=1)
        self.layer4 = self._make_layer(64, 64, 3, stride=1)
        #分类用的全连接
        self.fc = nn.Linear(256, num_classes)

    def _make_layer(self, inchannel, outchannel, block_num, stride=1):
        #构建layer,包含多个residual block
        shortcut = nn.Sequential(
            nn.Conv2d(inchannel, outchannel, 1, stride, bias=False),
            nn.BatchNorm2d(outchannel))
        layer = []
        layer.append(ResidualBloak(inchannel, outchannel, stride, shortcut))
        for i in range(1, block_num):
            layer.append(ResidualBloak(outchannel, outchannel))
        return nn.Sequential(*layer)

    def forward(self, x):
        x = self.pre(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = F.avg_pool2d(x, 7)
        x = x.view(x.size(0), -1)
        return self.fc(x)

def getData(): # 定义对数据的预处理  
    transform = transforms.Compose([
        transforms.Resize(40),
        transforms.RandomHorizontalFlip(),
        transforms.RandomCrop(32),
        transforms.ToTensor()])
    #训练集
    trainset = tv.datasets.CIFAR10(root='/data/', train=True, transform=transform, download=True)
    trainset_loader = DataLoader(trainset, batch_size=4, shuffle=True)
    #测试集
    testset = tv.datasets.CIFAR10(root='/data/', train=False, transform=transform, download=True)
    testset_loader = DataLoader(testset, batch_size=4, shuffle=False)
    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    return trainset_loader, testset_loader, classes

def train(): #训练
    trainset_loader, testset_loader, _ = getData() #获取数据
    net = ResNet34(10)
    print(net)
    criterion = nn.CrossEntropyLoss()
    optimizer = t.optim.SGD(net.parameters(), lr=0.001, momentum=0.9) #优化器

    for epoch in range(1):
        for step, (inputs,labels) in enumerate(trainset_loader):
            optimizer.zero_grad() #梯度清零
            output = net(inputs)
            loss = criterion(output, labels)
            loss.backward()
            optimizer.step()
            if step % 10 ==9:
                acc = test(net, testset_loader)
                print('Epoch', epoch, '|step ', step, 'loss: %.4f' %loss.item(), 'test accuracy:%.4f' %acc)
    print('Finished Training')
    return net

def test(net, testdata): #测试集
    correct, total = .0, .0
    for inputs, label in testdata:
        net.eval()
        output = net(inputs)
        _, predicted = t.max(output, 1) #分类结果
        total += label.size(0)
        correct += (predicted == label).sum()
    return float(correct) / total

if __name__ == '__main__':
    net = train()