(Pytorch)ResNet代码复现CIFAR-10数据集

resnet.py

'''ResNet-18 Image classfication for cifar-10 with PyTorch 

'''
import torch
import torch.nn as nn
import torch.nn.functional as F


class ResidualBlock(nn.Module):
    def __init__(self, inchannel, outchannel, stride=1):
        super(ResidualBlock, self).__init__()
        self.left = nn.Sequential(
            nn.Conv2d(inchannel, outchannel, kernel_size=3, stride=stride, padding=1, bias=False),
            nn.BatchNorm2d(outchannel),
            nn.ReLU(inplace=True),
            nn.Conv2d(outchannel, outchannel, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(outchannel)
        )
        self.shortcut = nn.Sequential()
        if stride != 1 or inchannel != outchannel:
            self.shortcut = nn.Sequential(
                nn.Conv2d(inchannel, outchannel, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(outchannel)
            )

    def forward(self, x):
        out = self.left(x)
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class ResNet(nn.Module):
    def __init__(self, ResidualBlock, num_classes=10):
        super(ResNet, self).__init__()
        self.inchannel = 64
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(),
        )
        self.layer1 = self.make_layer(ResidualBlock, 64, 2, stride=1)
        self.layer2 = self.make_layer(ResidualBlock, 128, 2, stride=2)
        self.layer3 = self.make_layer(ResidualBlock, 256, 2, stride=2)
        self.layer4 = self.make_layer(ResidualBlock, 512, 2, stride=2)
        self.fc = nn.Linear(512, num_classes)

    def make_layer(self, block, channels, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)  # strides=[1,1]
        layers = []
        for stride in strides:
            layers.append(block(self.inchannel, channels, stride))
            self.inchannel = channels
        return nn.Sequential(*layers)

    def forward(self, x):
        out = 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.fc(out)
        return out


def ResNet18():
    return ResNet(ResidualBlock)

train.py

import os
import sys
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

from resnet import ResNet18

# 定义是否使用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 超参数设置
epochs = 300  # 遍历数据集次数
BATCH_SIZE = 128  # 批处理尺寸(batch_size)
LR = 0.1  # 学习率

# 准备数据集并预处理
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),  # 先四周填充0，在吧图像随机裁剪成32*32
    transforms.RandomHorizontalFlip(),  # 图像一半的概率翻转，一半的概率不翻转
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),  # R,G,B每层的归一化用到的均值和方差
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True,
                                             transform=transform_train)  # 训练数据集
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True,
                                           num_workers=0)  # 生成一个个batch进行批训练，组成batch的时候顺序打乱取

val_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform_test)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=100, shuffle=False, num_workers=0)
val_num = len(val_dataset)
train_num = len(train_dataset)
train_steps = len(train_loader)
val_steps = len(val_loader)
# Cifar-10的标签
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# 模型定义-ResNet
model = ResNet18().to(device)

# 定义损失函数和优化方式
loss_function = nn.CrossEntropyLoss()  # 损失函数为交叉熵，多用于多分类问题
optimizer = optim.SGD(model.parameters(), lr=LR, momentum=0.9,
                      weight_decay=5e-4)  # 优化方式为mini-batch momentum-SGD，并采用L2正则化（权重衰减）
resume = True  # 设置是否需要从上次的状态继续训练
if resume:
    if os.path.isfile("ResNet.pth"):
        print("Resume from checkpoint...")
        checkpoint = torch.load("ResNet.pth")
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        initepoch = checkpoint['epoch'] + 2
        print("====>loaded checkpoint (epoch{})".format(checkpoint['epoch'] + 1))
    else:
        print("====>no checkpoint found.")
        initepoch = 1  # 如果没进行训练过，初始训练epoch值为1

writer = SummaryWriter("logs")
# 训练
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[136, 185], gamma=0.1)
for epoch in range(initepoch - 1, epochs):
    # train
    print("-------第 {} 轮训练开始-------".format(epoch + 1))
    model.train()
    train_acc = 0.0
    running_loss = 0.0
    train_bar = tqdm(train_loader, file=sys.stdout)
    for step, data in enumerate(train_bar):
        images, labels = data
        optimizer.zero_grad()
        outputs = model(images.to(device))
        loss = loss_function(outputs, labels.to(device))
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

        train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1, epochs, loss)
        _, predict = torch.max(outputs, dim=1)
        train_acc += torch.eq(predict, labels.to(device)).sum().item()
    train_loss = running_loss / train_steps
    train_accurate = train_acc / train_num
    # val
    model.eval()
    val_acc = 0.0
    running_loss = 0.0
    with torch.no_grad():
        val_bar = tqdm(val_loader, file=sys.stdout)
        for step, val_data in enumerate(val_bar):
            val_images, val_labels = val_data
            outputs = model(val_images.to(device))
            loss = loss_function(outputs, val_labels.to(device))
            running_loss += loss.item()
            _, predict = torch.max(outputs, dim=1)
            val_acc += torch.eq(predict, val_labels.to(device)).sum().item()
    val_loss = running_loss / val_steps
    val_accurate = val_acc / val_num
    scheduler.step()
    print(optimizer.state_dict()['param_groups'][0]['lr'])
    print('[epoch %d] train_loss: %.3f val_loss:%.3f train_accuracy:%.3f val_accuracy: %.3f' %
          (epoch + 1, train_loss, val_loss, train_accurate, val_accurate))
    writer.add_scalars('loss',
                       {'train': train_loss, 'val': val_loss}, global_step=epoch)
    writer.add_scalars('acc',
                       {'train': train_accurate, 'val': val_accurate}, global_step=epoch)
    # 保存断点
    checkpoint = {"model_state_dict": model.state_dict(),
                  "optimizer_state_dict": optimizer.state_dict(),
                  "epoch": epoch}
    path_checkpoint = "ResNet.pth"
    torch.save(checkpoint, path_checkpoint)
    print("保存模型成功")
print('Finished Training')
writer.close()

程序设置了断点续训，可以接着训练，查看日志可以用tensorboard