机器学习的API
机器学习模型分为四步:
1.定义模型、定义损失函数、定义损失函数优化器
2.前向传播->y’
3.计算损失->y-y’
4.反向传播->参数更新
# 一、定义model 、loss 、optimizer
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
if torch.cuda.is_available():
print("CUDA is enable!")
model = model.cuda()
model.train()
# 开始训练
for epoch in range(epoches_num):
print('*' * 40)
running_loss = 0.0
running_acc = 0.0
# 训练
for i, data in enumerate(train_loader, 1):
img, label = data
# 拥有GPU的小伙伴还是推荐使用GPU训练
if torch.cuda.is_available():
img = Variable(img).cuda()
label = Variable(label).cuda()
else:
img = Variable(img)
label = Variable(label)
# 二、前向传播
out = model(img)
# 三、计算损失
loss = criterion(out, label)
running_loss += loss.item() * label.size(0)
_, pred = torch.max(out, 1)
num_correct = (pred == label).sum()
accuracy = (pred == label).float().mean()
running_acc += num_correct.item()
# 四、反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Finish {} Loss: {:.6f}, Acc: {:.6f}'.format(epoch + 1, running_loss / (len(train_dataset)),
running_acc / (len(train_dataset))))
# 保存模型
torch.save(model, 'cnn.pt')