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，并且还可以进行网络模型的保存
