pytorch学习：梯度下降代码

# -*- coding: utf-8 -*-
"""
Created on Sun Sep  2 15:54:06 2018

@author: www
"""

import numpy as np
import torch
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
from torch import nn
from torch.autograd import Variable
import time
import matplotlib.pyplot as plt

def get_data(x):
     x = np.array(x, dtype='float32')/255
     x = (x-0.5) / 0.5  #标准化
     x = x.reshape((-1,)) #拉平
     x = torch.from_numpy(x)
     return x

train_set = MNIST('./data', train=True, transform=get_data, download=True) # 载入数据集，申明定义的数据变换
test_set = MNIST('./data', train=False, transform=get_data, download=True)     

#定义Loss函数
criterion = nn.CrossEntropyLoss()

#自定义随机梯度下降
def sgd_update(parameters, lr):
     for param in parameters:
          param.data = param.data - lr * param.grad.data
          
#现将batch_size设置为1.看看有什么效果
train_data = DataLoader(train_set, batch_size=1, shuffle=True)
#使用Sequential定义神经网络
net = nn.Sequential(
     nn.Linear(784,200),
     nn.ReLU(),
     nn.Linear(200, 10),
)
#开始训练
losses1 = []
idx = 0
start = time.time()
for e in range(5):
     train_loss = 0
     for im, label in train_data:
          im = Variable(im)
          label = Variable(label)
          #前向传播
          out = net(im)
          loss = criterion(out, label)
          #反向传播
          net.zero_grad()
          loss.backward()
          sgd_update(net.parameters(), 1e-2)
          #记录误差
          train_loss += loss.item()
          if idx%30 == 0:
               losses1.append(loss.item())
          idx += 1
     print('epoch:{},Train Loss:{}'.format(e, train_loss/len(train_data)))
end = time.time()
print('Time={}'.format(end - start))

x_axis = np.linspace(0, 5, len(losses1), endpoint=True)
plt.semilogy(x_axis, losses1, label='batch_size=1')
plt.legend(loc='best')
#运行后可以看到，loss 在剧烈震荡，因为每次都是只对一个样本点做计算，每一层的梯度都具有很高的
#随机性，而且需要耗费了大量的时间    

#接下来吧batch_size改成64
train_data = DataLoader(train_set, batch_size=64, shuffle=True)
# 使用 Sequential 定义 3 层神经网络
net = nn.Sequential(
    nn.Linear(784, 200),
    nn.ReLU(),
    nn.Linear(200, 10),
)

# 开始训练
losses2 = []
idx = 0
start = time.time() # 记时开始
for e in range(5):
    train_loss = 0
    for im, label in train_data:
        im = Variable(im)
        label = Variable(label)
        # 前向传播
        out = net(im)
        loss = criterion(out, label)
        # 反向传播
        net.zero_grad()
        loss.backward()
        sgd_update(net.parameters(), 1e-2)
        # 记录误差
        train_loss += loss.data[0]
        if idx % 30 == 0:
            losses2.append(loss.data[0])
        idx += 1
    print('epoch: {}, Train Loss: {:.6f}'
          .format(e, train_loss / len(train_data)))
end = time.time() # 计时结束
print('使用时间: {:.5f} s'.format(end - start))      
          
x_axis = np.linspace(0, 5, len(losses2), endpoint=True)
plt.semilogy(x_axis, losses2, label='batch_size=64')
plt.legend(loc='best')

#通过上面的结果可以看到 loss 没有 batch 等于 1 震荡那么距离，
#同时也可以降到一定的程度了，时间上也比之前快了非常多，因为按照 
#batch 的数据量计算上更快，同时梯度对比于 batch size = 1 
#的情况也跟接近真实的梯度，所以 batch size 的值越大，梯度也
#就越稳定，而 batch size 越小，梯度具有越高的随机性，这里 
#batch size 为 64，可以看到 loss 仍然存在震荡，但这并没有
#关系，如果 batch size 太大，对于内存的需求就更高，同时也不
#利于网络跳出局部极小点，所以现在普遍使用基于 batch 的随机梯
#度下降法，而 batch 的多少基于实际情况进行考虑


train_data = DataLoader(train_set, batch_size=64, shuffle=True)
# 使用 Sequential 定义 3 层神经网络
net = nn.Sequential(
    nn.Linear(784, 200),
    nn.ReLU(),
    nn.Linear(200, 10),
)

optimzier = torch.optim.SGD(net.parameters(), 1e-2)
# 开始训练

#下面我们来使用 pytorch 自带的优化器来实现随机梯度下降
train_data = DataLoader(train_set, batch_size=64, shuffle=True)
# 使用 Sequential 定义 3 层神经网络
net = nn.Sequential(
    nn.Linear(784, 200),
    nn.ReLU(),
    nn.Linear(200, 10),
)

optimzier = torch.optim.SGD(net.parameters(), 1e-2)
# 开始训练
start = time.time() # 记时开始
for e in range(5):
    train_loss = 0
    for im, label in train_data:
        im = Variable(im)
        label = Variable(label)
        # 前向传播
        out = net(im)
        loss = criterion(out, label)
        # 反向传播
        optimzier.zero_grad()
        loss.backward()
        optimzier.step()
        # 记录误差
        train_loss += loss.item()
    print('epoch: {}, Train Loss: {:.6f}'
          .format(e, train_loss / len(train_data)))
end = time.time() # 计时结束
print('使用时间: {:.5f} s'.format(end - start))