【PyTorch】PyTorch入门教程六

feedforward_neural_network

CPU版本

前面简单介绍了用PyTorch进行逻辑回归和线性回归，在逻辑回归中使用一个线性层对mnist数据集进行了分类，其实就是一个简单的神经网络，模子都是一样的，在这一节中我们使用两个线性层来看看如何构造。

首先还是从mnist数据集中加载数据。

import torch
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from torch.autograd import Variable


# Hyper Parameters 
input_size = 784
hidden_size = 500
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 0.001

# MNIST Dataset 
train_dataset = dsets.MNIST(root='./data', 
                            train=True, 
                            transform=transforms.ToTensor(),  
                            download=True)

test_dataset = dsets.MNIST(root='./dat