1、卷积层编程要点

画圈参数需要仔细琢磨。

2、代码

完整的MNIST数据集5层神经网络代码，CPU跑用时221.5秒，GPU加速用时19秒。

import torch
import torchvision as tv
import time

"导入数据"


def load_data(batch=500):
    train_set = torch.utils.data.DataLoader(
        tv.datasets.MNIST("../data", train=True, download=True,
                          transform=tv.transforms.Compose(
                              [tv.transforms.ToTensor(), tv.transforms.Normalize((0.1307,), (0.3081,))])),
        batch_size=batch, shuffle=True
    )

    test_set = torch.utils.data.DataLoader(
        tv.datasets.MNIST("../data", train=False, download=False,
                          transform=tv.transforms.Compose(
                              [tv.transforms.ToTensor(), tv.transforms.Normalize((0.1307,), (0.3081,))])),
        batch_size=batch, shuffle=True
    )
    return train_set, test_set


"建立网络"


class net(torch.nn.Module):
    def __init__(self):
        super(net, self).__init__()
        self.conv0 = torch.nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5)
        self.pooling1 = torch.nn.MaxPool2d(kernel_size=2)
        self.conv2 = torch.nn.Conv2d(in_channels=10, out_channels=20, kernel_size=5)
        self.pooling3 = self.pooling1
        self.fc4 = torch.nn.Linear(320, 10)

    def forward(self, x):
        batch = x.size(0)
        x = self.pooling1(torch.nn.functional.relu(self.conv0(x)))
        x = self.pooling3(torch.nn.functional.relu(self.conv2(x)))
        x = self.fc4(x.view([batch, -1]))
        return x


def train(epochs=1):
    for epoch in range(epochs):
        for imgs, labels in train_set:
            imgs, labels = imgs.to(device), labels.to(device)
            for i in range(100):
                y = model(imgs)
                loss = criterion(y, labels)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()


def test():
    total, correct = 0, 0
    with torch.no_grad():
        for imgs, labels in test_set:
            imgs, labels = imgs.to(device), labels.to(device)
            y = model(imgs)
            _, predict = torch.max(y.data, dim=1)
            total += predict.size(0)
            correct += (predict == labels).sum().item()
        return correct / total


if __name__ == "__main__":
    begin_time = time.time()
    # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    device = torch.device("cpu")
    test_set, train_set = load_data()
    model = net()
    model = model.to(device)

    epochs = 10
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
    train()
    rate = test()
    print(rate)
    end_time = time.time()
    print("用时: %f 秒" % (end_time - begin_time))