课程请见 《PyTorch深度学习实践》
# PyTorch
import torch
from torch import nn
from torch import optim
from torch.utils import data
class LinearNet(nn.Module):
def __init__(self, input_dim, output_dim=1):
super(LinearNet, self).__init__()
self.linear = nn.Linear(input_dim, output_dim)
def forward(self, x):
return self.linear(x)
def init_weight(m):
if type(m) == nn.Linear:
nn.init.normal(m.weight, std=0.01)
net = LinearNet(1)
net.apply(init_weight)
x_data = torch.Tensor([[1.0], [2.0], [3.0]])
y_data = torch.Tensor([[2.0], [4.0], [6.0]])
criterion = nn.MSELoss()
optimizer = optim.ASGD(net.parameters(), lr=0.01)
num_epochs = 1000
for epoch in range(num_epochs):
_y = net(x_data)
loss = criterion(_y, y_data)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('第{}轮, loss为{:.8f}'.format(epoch+1, float(loss.data)))