Pytorch_3:可执行的最小程序

from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms, models

def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    # 0. Settings
    torch.manual_seed(100)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    batch_size = 32
    epochs = 20
    lr = 0.01
    momentum = 0.90

    # 1. Dataset
    dataset_train = datasets.CIFAR10('../data', train=True, download=True,
                                     transform=transforms.Compose([
                                         transforms.ToTensor(),
                                         transforms.Normalize((0.1307,), (0.3081,))
                                     ]))
    dataset_test = datasets.CIFAR10('../data', train=False, download=True,
                                    transform=transforms.Compose([
                                        transforms.ToTensor(),
                                        transforms.Normalize((0.1307,), (0.3081,))
                                    ]))
    train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=4)
    test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=batch_size, shuffle=True, num_workers=4)

    # 2. Model
    model = models.resnet18(pretrained=True, progress=True)  # 从模型库中调用ResNet18
    print(model)
    # num_ftrs = model.fc.in_features
    # model.fc = nn.Linear(num_ftrs, 10)  # 得到新的模型，更改fc层
    # model = model.to(device)
    # print(model)

    # 3. Optimizer
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)

    # 4. Training
    for epoch in range(epochs):
        train(model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        # save the current model
        if (epoch + 1) % 10 == 0:
            torch.save(model, "train_cnn_%3d.pt" % (epoch + 1))


if __name__ == '__main__':
    main()