一、须知
1.本代码所用数据集为CIFAR10,可通过以下代码段进行下载并加载
需要引用 import torchvision
train_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=True, transform=torchvision.transforms.ToTensor())
test_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=False, transform=torchvision.transforms.ToTensor())
2.网络不支持数据集中各图片尺寸相互不一的情况,若自行构建数据集或加载别的数据集,请先对数据集尺寸做成统一格式。推荐更改为3 * 32 * 32 ,若更改为其他格式,自行计算nn.Flatten()之后的像素总数,并替换掉nn.Linear(1024, 10)中的1024
3.测试集结果不输出最终类别判断,仅支持正确率(正确个数/测试集总个数) 的输出
4.支持tensorboard
5.加入每100次迭代计算时间差
6.未加入激活函数,需要自行添加
7.由于基础框架比较简单,模型表现效果略差,运行165epoch时,测试集取得最高准确率68.5%
二、网络模型框架
基本构架思路为
读取数据→构建minibacth→选择GPU或CPU训练→选择损失函数→构建前向传递网络→选择GSD模型进行下降并设置超参→开始迭代→计算损失函数→反向传播→更新参数→输出结果→测试
三、完整代码
可直接在kaggle的code上运行,数据集选择cifar10-python即可
import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time
train_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=True, transform=torchvision.transforms.ToTensor())
test_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=False, transform=torchvision.transforms.ToTensor())
train_dataloader = DataLoader(train_data, batch_size=64, drop_last=True)
test_dataloader = DataLoader(test_data, batch_size=64, drop_last=True)
# print(len(train_dataloader)) #781
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
test_data_size = len(test_dataloader) * 64
print(f'测试集大小为:{test_data_size}')
writer = SummaryWriter("../model_logs")
loss_fn = nn.CrossEntropyLoss(reduction='mean')
loss_fn = loss_fn.to(device)
time_able = False # True
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.model1 = nn.Sequential(
nn.Conv2d(3, 32, 5, padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, padding=2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(1024, 64),# 182528
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model1(x)
return x
model = Model()
model = model.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
epoch = 50
running_loss = 0
total_train_step = 0
total_test_step = 0
if time_able:
str_time = time.time()
for i in range(epoch):
print(f'第{i + 1}次epoch')
for data in train_dataloader:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
output = model(imgs)
loss = loss_fn(output, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step += 1
if total_train_step % 100 == 0:
if time_able:
end_time = time.time()
print(f'{str_time-end_time}')
print(f'第{total_train_step}次训练,loss = {loss.item()}')
writer.add_scalar("train_loss", loss.item(), total_train_step)
# 测试
total_test_loss = 0
total_accuracy = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
outputs = model(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss
accuracy = (outputs.argmax(1) == targets).sum()
total_accuracy += accuracy
total_test_loss = total_test_loss / test_data_size
print(f'整体测试集上的loss = {total_test_loss}')
print(f'整体测试集正确率 = {total_accuracy / test_data_size}')
writer.add_scalar("test_loss", total_test_loss.item(), total_test_step)
writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)
total_test_step += 1
writer.close()