梯度下降公式很简单
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
#随机梯度下降法,从0开始自己实现
def data_tf(x):
x = np.array(x,dtype='float32')/255 #讲数据变成0-1之间
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)
#定义loss函数
criterion = nn.CrossEntropyLoss()
#定义梯度下降公式:lr学习率,parameters参数
def sgd_update(parameters,lr):
for param in parameters:
param.data = param.data - lr * param.grad.data
train_data = DataLoader(train_set,batch_size=64,shuffle=True)
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(),0.01)
#记录误差
train_loss += loss.data
if idx%30 == 0:
losses1.append(loss.data)
idx += 1
print('epoch:{},Train loss:{}'.format(e,train_loss/len(train_data)))
end = time.time()
print('使用时间:{}'.format(end-start))
x_axis = np.linspace(0,5,len(losses1),endpoint=True)
plt.semilogy(x_axis,losses1,label='batch_size = 1')
plt.show()
batch_size看电脑配置,越小越不稳定
PyTorch内置随机梯度下降法
train_data = DataLoader(train_set, batch_size=64, shuffle=True)
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.data
print('epoch: {}, Train Loss: {:.6f}'
.format(e, train_loss / len(train_data)))
end = time.time() #
print('时间: {} s'.format(end - start))