基于 PyTorch 的 cifar-10 图像分类

import torch import torch.nn as nn import torchvision from vggnet import VGGNet from load_cifar10 import test_loader, train_loader import os import tensorboardX # 查看gpu device = torch.device("cuda" if torch.cuda.is_available() else "cpu") epoch_num = 100 lr = 0.001 net = VGGNet().to(device) # 多分类问题交叉商 loss_func = nn.CrossEntropyLoss() # 优化器 optimizer = torch.optim.Adam(net.parameters(), lr=lr) # 采用变长学习率，每五个epoch指数衰减 scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.9) # 用于记录过程的log if not os.path.exists("log"): os.mkdir("log") writer = tensorboardX.SummaryWriter("log") # step计数器 step_n = 0 for epoch in range(epoch_num): print("epoch is", epoch) net.train() # train BN dropout batchsize = 128 for i, data in enumerate(train_loader): inputs, labels = data inputs, labels = inputs.to(device), labels.to(device) outputs = net(inputs) # 获取输出 # 计算loss loss = loss_func(outputs, labels) optimizer.zero_grad() # 优化qi梯度为 loss.backward() # 反向传播 optimizer.step() # 更新参数 # 打印batch的loss # print("step", i, "loss is :", loss.item()) _, pred = torch.max(outputs.data, dim=1) correct = pred.eq(labels.data).cpu().sum() print("train step", i, "loss is :", loss.item(), "batch correct is:", 100.0 * correct / batchsize) # 记录loss。correct writer.add_scalar("train loss", loss.item(), global_step=step_n) writer.add_scalar("train correct", 100.0 * correct / batchsize, global_step=step_n) im = torchvision.utils.make_grid(inputs) writer.add_image("train image", im, global_step=step_n) step_n += 1 # 保存每个epoch模型 if not os.path.exists("models"): os.mkdir("models") torch.save(net.state_dict(), "models/{}.pth".format(epoch + 1)) # 学习率更新 scheduler.step() print("lr is:", optimizer.state_dict()["param_groups"][0]["lr"]) # 测试 sum_loss = 0 sum_correct = 0 net.eval() for i, data in enumerate(test_loader): inputs, labels = data inputs, labels = inputs.to(device), labels.to(device) outputs = net(inputs) # 获取输出 # 计算loss loss = loss_func(outputs, labels) # print("step", i, "loss is :", loss.item()) _, pred = torch.max(outputs.data, dim=1) correct = pred.eq(labels.data).cpu().sum() sum_correct += correct.item() sum_loss += loss.item() # 记录loss。correct # print("test step", i, "test loss is :", loss.item(), "test batch correct is:", 100.0 * correct / batchsize) im = torchvision.utils.make_grid(inputs) writer.add_image("train image", im, global_step=step_n) test_loss = sum_loss * 1.0 / len(test_loader) test_correct = sum_correct * 100.0 / len(test_loader) / batchsize writer.add_scalar("test loss", loss.item(), global_step=epoch + 1) writer.add_scalar("test correct", 100.0 * correct / batchsize, global_step=epoch + 1) print("test epoch", epoch + 1, "test loss is :", test_loss, "test batch correct is:", test_correct) writer.close()