1,准备数据集(以CIFAR10为例)
train_data = torchvision.datasets.CIFAR10("../data", train=True, transform=torchvision.transforms.ToTensor(), download=True) test_data = torchvision.datasets.CIFAR10("../data", train=False, transform=torchvision.transforms.ToTensor(), download=True)
(补充)计算数据集的长度用len()
train_data_size = len(train_data) test_data_size = len(test_data) print("训练集的长度为:".format(train_data_size)) print("测试集的长度为:".format(test_data_size))
2,利用dataloader来加载数据集
train_dataloader = DataLoader(train_data, batch_size=64) test_dataloader = DataLoader(test_data, batch_size=64)
3,创建网络模型
class Sjwl(nn.Module): def __init__(self): super(Sjwl, self).__init__() self.model = nn.Sequential( nn.Conv2d(3, 32, 5, 1, 2), nn.MaxPool2d(2), nn.Conv2d(32, 32, 5, 1, 2), nn.MaxPool2d(2), nn.Conv2d(32, 64, 5, 1, 2), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(64 * 4 * 4, 64), nn.Linear(64, 10) ) def forward(self, x): x = self.model(x) return x 4,使用损失函数,这里使用的是交叉熵损失函数
loss_fountion = nn.CrossEntropyLoss()
5,使用优化器
learnspeed = 0.01 optim = torch.optim.SGD(sjwl.parameters(), lr=learnspeed)
6,设置训练网络的一些参数
# 记录训练的次数 total_train_step = 0 # 记录测试的次数 total_test_step = 0 # 训练的轮数 epoch = 10 #添加tensorboard writer=SummaryWriter("../logs_train")
7,训练步骤的开始
for i in range(epoch): print("-----第{}轮训练开始------".format(i + 1)) # 训练步骤的开始 for data in train_dataloader: imgs, targets = data output = sjwl(imgs) loss = loss_fountion(output, targets) # 优化器优化模型 optim.zero_grad() loss.backward() optim.step() total_train_step = total_train_step + 1 if total_train_step % 100 == 0: print("训练次数为{}时,Loss:{}".format(total_train_step, loss.item())) writer.add_scalar("train_loss",loss.item(),total_train_step)
8,测试步骤的开始
# 测试步骤的开始 total_test_loss = 0 total_accuracy=0 with torch.no_grad(): for data in test_dataloader: imgs, targets = data output = sjwl(imgs) loss = loss_fountion(output, targets) total_test_loss = total_test_loss + loss.item() #计算正确率的次数 accuracy=(output.argmax(1)==targets).sum() total_accuracy=total_accuracy+accuracy print("整体测试集上的loss:{}".format(total_test_loss)) print("整体测试集上的正确率:{}".format(accuracy)) writer.add_scalar("test_loss",total_test_loss,total_test_step) writer.add_scalar("test_accuracy",accuracy,total_test_step ) total_test_step=total_test_step+1 torch.save(sjwl,"sjwl_{}.pth".format(i)) print("模型已保存") writer.close()
其中还使用了tensorboard,并且还可以进行网络模型的保存