深入理解ResNet18代码

发布时间:2023-05-20

ResNet18介绍

ResNet是一种非常有名的卷积神经网络,具有非常深的层数。当网络的层数加深时,由于梯度消失的问题,训练时就会变得非常困难。而ResNet通过引入跨层连接的方法解决了这个问题。ResNet18是ResNet的一种常用的轻量级模型,本文将介绍ResNet18的详细代码。

ResNet18结构

ResNet18由多个卷积块组成,每个卷积块包含不同数量的卷积层和跨层连接。在ResNet18中,存在4个卷积块,每个卷积块有不同数量的卷积层,具体为2、2、2、2个卷积层。

class BasicBlock(nn.Module):
    expansion = 1
    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion*planes)
            )
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out

卷积块

卷积块是ResNet18的核心组成部分,包含多个卷积层和跨层连接。每个卷积块的第一个卷积层的输入通道数和输出通道数不一样,需要进行维度变换。因此,需要在卷积块中添加一条shortcut路径,用于跨越多层进行连接。在代码实现中,shortcut路径基于该卷积块输入和输出的尺寸和通道数,选择是采用1x1卷积,还是直接新建一个卷积层进行维度变换。

ResNet18模型代码

ResNet18模型代码由多个卷积块组成,每个卷积块内部包含多个卷积层和跨层连接。ResNet18模型还包括一个输入层和一个全连接层。在输入层通过一系列卷积层和池化层微调数据,最后经过全连接层输出结果。

class ResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10):
        super(ResNet, self).__init__()
        self.in_planes = 64
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
        self.linear = nn.Linear(512*block.expansion, num_classes)
    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out

模型训练代码

模型训练代码包括数据的加载、模型的训练与验证、损失函数的计算和梯度反向传播。其中,需要注意的是,交叉熵损失函数和SGD优化器的参数设置。

device = 'cuda' if torch.cuda.is_available() else 'cpu'
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
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=2)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)
net = ResNet(BasicBlock, [2, 2, 2, 2]).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
for epoch in range(200):
    net.train()
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
    net.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(testloader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)
            loss = criterion(outputs, targets)
            test_loss += loss.item()
            _, predicted = outputs.max(1)
            correct += predicted.eq(targets).sum().item()
    test_loss /= len(testloader.dataset)
    print('Epoch: {}, Test Loss: {:.4f}, Test Acc: {:.2f}%'.format(
        epoch+1, test_loss, 100.*correct/len(testloader.dataset)))