ResNet18训练cifar-10数据集

在首先进行了数据增广和数据标准化。

train_augs = torchvision.transforms.Compose([torchvision.transforms.Resize(224),
                                             torchvision.transforms.RandomHorizontalFlip(),
                                             #color_aug, shape_aug,
                                             torchvision.transforms.ToTensor(),
                                             torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                                                                              std=[0.229, 0.224, 0.225])])

每隔2个epoch就将学习乘以0.9

StepLR = torch.optim.lr_scheduler.StepLR(optimizer,step_size=2, gamma=0.90)

在Dataloader函数让shuffle=True进行洗牌。



可以看见模型已经收敛，但是测试的精度只有85%，模型陷入了局部最优解。后续看看能不能进行优化。
训练代码如下：

import torch
from torch import nn
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torch.nn import functional as F

"""获取设备"""
def try_gpu(i = 0):
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')
device = try_gpu()
print(f'device:{device}')

class Residual(nn.Module):  #@save
    def __init__(self, input_channels, num_channels,
                 use_1x1conv=False, strides=1):
        super().__init__()
        self.conv1 = nn.Conv2d(input_channels, num_channels,
                               kernel_size=3, padding=1, stride=strides)
        self.conv2 = nn.Conv2d(num_channels, num_channels,
                               kernel_size=3, padding=1)
        if use_1x1conv:
            self.conv3 = nn.Conv2d(input_channels, num_channels,
                                   kernel_size=1, stride=strides)
        else:
            self.conv3 = None
        self.bn1 = nn.BatchNorm2d(num_channels)
        self.bn2 = nn.BatchNorm2d(num_channels)

    def forward(self, X):
        Y = F.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
        if self.conv3:
            X = self.conv3(X)
        Y += X
        return F.relu(Y)

b1 = nn.Sequential(nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
                  nn.BatchNorm2d(64), nn.ReLU(),
                  nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
def resnet_block(input_channels, num_channels, num_residuals,
                first_block=False):
    blk = []
    for i in range(num_residuals):
        if i == 0 and not first_block:
            blk.append(Residual(input_channels, num_channels,
                                use_1x1conv=True, strides=2))
        else:
             blk.append(Residual(num_channels, num_channels))
    return blk

b2 = nn.Sequential(*resnet_block(64, 64, 2, first_block=True))
b3 = nn.Sequential(*resnet_block(64, 128, 2))
b4 = nn.Sequential(*resnet_block(128, 256, 2))
b5 = nn.Sequential(*resnet_block(256, 512, 2))

net = nn.Sequential(b1, b2, b3, b4, b5,
                    nn.AdaptiveAvgPool2d((1,1)),
                    nn.Flatten(), nn.Linear(512, 10))

# X = torch.rand(size=(1, 3, 224, 224))
# for layer in net:
#     X = layer(X)
#     print(layer.__class__.__name__,'output shape:\t', X.shape)
net = net.to(device)
def initial(layer):
    if isinstance(layer, nn.Linear) or isinstance(layer, nn.Conv2d):
        nn.init.xavier_normal_(layer.weight.data)
net.apply(initial)

"""定义loss"""
loss = nn.CrossEntropyLoss()
loss = loss.to(device)

"""定义优化器"""
lr = 0.01
optimizer = torch.optim.SGD(net.parameters(), lr,momentum=0.9, weight_decay=0.0005)
StepLR = torch.optim.lr_scheduler.StepLR(optimizer,step_size=2, gamma=0.90)

"""加载数据"""
color_aug = torchvision.transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5)
shape_aug = torchvision.transforms.RandomResizedCrop((200, 200), scale=(0.1, 1), ratio=(0.5, 2))

train_augs = torchvision.transforms.Compose([torchvision.transforms.Resize(224),
                                             torchvision.transforms.RandomHorizontalFlip(),
                                             #color_aug, shape_aug,
                                             torchvision.transforms.ToTensor(),
                                             torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                                                              std=[0.229, 0.224, 0.225])])
test_augs = torchvision.transforms.Compose([torchvision.transforms.Resize(224),
                                            torchvision.transforms.ToTensor(),
                                            torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

train_data = torchvision.datasets.CIFAR10(root='/root/dataset', train=True, download=False,
                                          transform=train_augs)
test_data = torchvision.datasets.CIFAR10(root='/root/dataset', train=False, download=False,
                                          transform=test_augs)

train_dataloader = DataLoader(dataset=train_data, batch_size=256, shuffle=True, drop_last=True)
test_dataloader = DataLoader(dataset=test_data, batch_size=256, shuffle=True, drop_last=True)

train_data_size = len(train_data)
test_data_size = len(test_data)

total_train_step = 0
total_test_step = 0
epoch = 10000


# 添加tensorboard
writer = SummaryWriter("/root/tf-logs")

max_acc = 0
for i in range(epoch):
    print("----------第 {} 轮训练开始----------".format(i + 1))

    net.train()
    total_train_loss = 0
    i = 0
    for data in train_dataloader:
        i += 1
        x, y = data
        x = x.to(device)
        y = y.to(device)

        y_hat = net(x)
        l = loss(y_hat, y)
        total_train_loss += l

        optimizer.zero_grad()
        l.backward()
        optimizer.step()
        total_train_step += 1
        print("训练次数: {}, Loss: {}".format(i, l.item()))

    StepLR.step()
    p = optimizer.param_groups[0]
    print('lr:{}'.format(p['lr']))

    net.eval()
    total_test_loss = 0
    total_accuracy = 0

    with torch.no_grad():
        for data in test_dataloader:
            x, y = data
            x = x.to(device)
            y = y.to(device)

            y_hat = net(x)
            l = loss(y_hat, y)
            total_test_loss += l

            accuracy = (y_hat.argmax(1) == y).sum()
            total_accuracy += accuracy

    print("整体训练集上的Loss: {}".format(total_train_loss))
    print("整体测试集上的Loss: {}".format(total_test_loss))
    print('整体测试集上的正确率: {}'.format(total_accuracy / test_data_size))
    writer.add_scalar("train_loss", total_train_loss, total_train_step)
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)
    total_test_step += 1

    if max_acc > 0.82 and total_accuracy / test_data_size > max_acc:
        torch.save(net, "vgg_net_accuracy_{}_{}.pth".format(
            total_accuracy / test_data_size, i + 1))
        print("模型已保存")
    max_acc = max(total_accuracy / test_data_size, max_acc)

writer.close()