【PyTorch】:ResNet34實現
阿新 • • 發佈:2019-02-19
# 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()