- 给定输入输出
- 定义一个模型
- 定义损失函数(loss function)和优化函数(optimizer)
- 训练一个过程
一个简单的例子:
import torch
import torch.nn as nn
N, D_in, H, D_out = 64, 1000, 100, 10
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)
class TwoLayerNet(torch.nn.Module):
def __init__(self, D_in, H, D_out):
super(TwoLayerNet, self).__init__()
self.linear1 = nn.Linear(D_in, H, bias=False)
self.linear2 = nn.Linear(H, D_out, bias=False)
def forward(self, x):
y_pred = self.linear2(self.linear1(x).clamp(min=0))
return y_pred
# model = TwoLayerNet(D_in, H, D_out)
if torch.cuda.is_available():
x = x.cuda()
y = y.cuda()
model = TwoLayerNet(D_in, H, D_out).cuda()
else:
model = TwoLayerNet(D_in,H, D_out)
# 损失函数和优化函数
loss_fn = nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr = 1e-4)
for it in range(500):
# Forward pass
y_pred = model(x)
# 计算loss
loss = loss_fn(y_pred, y)
print(it, loss.item())
# backward pass
optimizer.zero_grad() # 避免出现梯度累积
loss.backward() # 反向传播
# updata model parameters
optimizer.step()
训练结果: