模型训练套路——以cifar10为例
1.加载数据集
dataset_transform_compose = transforms.Compose([transforms.ToTensor()])
train_set = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=dataset_transform_compose, download=True)
test_set = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=dataset_transform_compose, download=True)
print(len(train_set))
print("测试集的大小为:{}".format(len(test_set)))
format的用法:format(##),##替换{}
输出:
50000
测试集的大小为:10000
2 利用DataLoader加载数据集
train_dataloader = DataLoader(train_set, batch_size=64)
test_dataloader = DataLoader(test_set, batch_size=64)
3 搭建网络
class model(nn.Module):
def __init__(self):
super(model, self).__init__()
self.model1 = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=(5, 5), stride=(1, 1), padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, (5, 5), padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, (5, 5), padding=2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(1024, 64),
nn.Linear(64, 10)
)
def forward(self, input):
input = self.model1(input)
return input
4 创建网络模型
model1 = model()
5 定义损失函数
bloss = nn.CrossEntropyLoss()
6 定义优化器
opti = torch.optim.Adam(model1.parameters(), lr=0.01)
7 设置网络的一些参数
train_step = 0
test_step = 0
epoch = 10
for i in range(epoch):
print("----------第{}轮训练开始了----------".format(i+1))
##训练开始
for data in train_dataloader:
imgs, targets = data
output = model1(imgs)
loss = bloss(output, targets)
writer.add_scalar("训练集上的loss", loss, train_step)
#优化器优化步骤
opti.zero_grad()
loss.backward()
opti.step()
train_step = train_step + 1
if train_step % 100 ==0:
print("第{}步训练, 损失为: {}".format(train_step, loss))
##测试开始
test_loss = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
output = model1(imgs)
loss = bloss(output, targets)
test_loss = test_loss + loss
print("整个测试集上的loss为:{}".format(test_loss))
writer.add_scalar("测试集上的loss", test_loss, test_step)
test_step = test_step + 1
torch.save(model1, "第{}轮的模型.pth".format(i+1))