优化算法2--动量法

起始点a到b点的梯度下降记为n1
b点到c点的实际梯度由两方面组成
v (t) =γv(t-1) +α∇b
+前面是n1的γ倍，通常取0.9（为什么<1？因为要减小早期的梯度影响）
+后面α是学习率，∇b是b点理论上的梯度下降

有很多点之后，v(t-1)指的是之前所有步骤累加的动量和

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

def data_tf(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=data_tf, download=True)
test_set = MNIST('./data', train=False, transform=data_tf, download=True)

criterion = nn.CrossEntropyLoss()
train_data = DataLoader(train_set, batch_size=64, shuffle=True)

net = nn.Sequential(
    nn.Linear(784, 200),
    nn.ReLU(),
    nn.Linear(200, 10)
)

optimizer = torch.optim.SGD(net.parameters(),lr=1e-2,momentum=0.9)

losses = []
start = time.time()
idx = 0
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()
        optimizer.step()
        train_loss += loss.data
        if idx % 30 == 0:
            losses.append(loss.data)
        idx += 1

    print('epoch: {}, Train Loss: {:.6f}'.format(e, train_loss / len(train_data)))
end = time.time()
print('时间: {:.5f} s'.format(end - start))