手写神经网络：Pytorch版

最近为了对比神经网络在不同语言中的运行效率，于是手搓了一个简单模型在不同语言中的代码，在这儿记录一下。
基于MINIST手写数字数据集，采用3层全连接网络，由Pytorch、Python、C++、CUDA四个版本实现。

0. 实验结果
1. Pytorch版
2. Python版
   2.1.Numpy版
   2.2.纯Python版
3. C++版
4. CUDA版

Python3.6

Pytorch1.7

import torch
from torchvision import datasets, transforms
import os,time
import torchvision
import numpy as np
from torch.autograd import Variable

transform = transforms.Compose([transforms.ToTensor(),
                               transforms.Normalize(mean=[0.5],std=[0.5])])

data_train = datasets.MNIST(root = "./data/",
                            transform=transform,
                            train = True,
                            download = True)

data_test = datasets.MNIST(root="./data/",
                           transform = transform,
                           train = False)

data_loader_train = torch.utils.data.DataLoader(dataset=data_train,
                                                batch_size = 100,
                                                shuffle = True,
                                                 num_workers=0)

data_loader_test = torch.utils.data.DataLoader(dataset=data_test,
                                               batch_size = 100,
                                               shuffle = True,
                                                num_workers=0)

print(len(data_train))


class Model(torch.nn.Module):

    def __init__(self):
        super(Model, self).__init__()
        self.dense = torch.nn.Sequential(torch.nn.Linear(28*28, 1024),
                                         torch.nn.ReLU(),
                                         torch.nn.Linear(1024,1024),
                                         torch.nn.ReLU(),
                                         torch.nn.Linear(1024, 10),
                                         # torch.nn.Softmax()  #这里不用再加softmax  因为nn.CrossEntropyLoss中默认有
                                         )

    def forward(self, x):

        x=x.view(-1,28*28)
        x = self.dense(x)
        return x

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Model().to(device)

cost = torch.nn.CrossEntropyLoss()
# 优化器
optimizer = torch.optim.SGD(model.parameters(),lr=0.01)
n_epochs = 5

for epoch in range(n_epochs):
    running_loss = 0.0
    running_correct = 0
    print("Epoch {}/{}".format(epoch, n_epochs))
    print("-" * 50)
    t1=time.time()
    for data in data_loader_train:
        X_train, y_train = data
        X_train, y_train = Variable(X_train).to(device), Variable(y_train).to(device)
        outputs = model(X_train)
        _, pred = torch.max(outputs.data, 1)
        optimizer.zero_grad()
        loss = cost(outputs, y_train)

        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        running_correct += torch.sum(pred == y_train.data)
    testing_correct = 0
    # print('Testing----')
    for data in data_loader_test:
        X_test, y_test = data
        X_test, y_test = Variable(X_test).to(device), Variable(y_test).to(device)
        outputs = model(X_test)
        _, pred = torch.max(outputs.data, 1)
        testing_correct += torch.sum(pred == y_test.data)
    print("Loss is:{:.4f}, Train Accuracy is:{:.4f}%, Test Accuracy is:{:.4f}".format(running_loss / len(data_train),
                                                                                      100 * running_correct / len(
                                                                                          data_train),
                                                                                      100 * testing_correct / len(
                                                                                          data_test)))
    t2=time.time()
    print('Time:',t2-t1)
torch.save(model.state_dict(), "model_parameter.pkl")