机器学习的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')