1、导入环境
import torch
2、设定训练数据大小
batch_size=32
input_shape=5
out_putshape=10
3、启动GPU 进行训练
torch.set_default_tensor_type(‘torch.cuda.FloatTensor’)
4、生成随机训练数据
from torch.autograd import Variable
x=Variable(torch.randan(bitch_size,input_shape))
y=Variable(torch.randan(batch_size,output_shape),requires_grad=False)
5、实现简单神经网络
注意:使用.cuda()扩展来确保模型在Gpu上运行。
model=torch.nn.Sequential(
torch.nn.Linear(input_shape,32),
torch.nn.Linear(32,output_shape),
).cuda()
6、定义损失函数
loss_function=torch.nn.MSELoss()
7、开始训练
learning_rate=0.001
for i in range(10):
y_pred=model(x)
loss=loss_function(y_pred,y)
print(loss.data[0])
#零梯度
model.zero_grad()
loss.backward()
#更新权值
for param in model.parameters():
param.data-=learning_rate*param.grad.data
代码参考:
python 深度学习实战-75个有关神经网络建模、强化学习与迁移学习的解决方案