《PyTorch深度学习实践》线性回归

课程请见 《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)))