"""All training mechanics visible in one small regression script."""
import torch
from torch import nn
torch.manual_seed(11)
torch.set_num_threads(1)
x = torch.tensor([[-2.], [-1.], [0.], [1.], [2.]])
y = 2*x + 1
model = nn.Linear(1,1)
optimizer = torch.optim.SGD(model.parameters(),lr=.1)
for epoch in range(100):
model.train()
optimizer.zero_grad(set_to_none=True)
prediction = model(x)
loss = nn.functional.mse_loss(prediction,y)
loss.backward()
optimizer.step()
model.eval()
with torch.no_grad():
print('weight, bias:',model.weight.item(),model.bias.item())
print('prediction for 3:',model(torch.tensor([[3.]])).item())
assert abs(model.weight.item()-2) < 1e-4
assert abs(model.bias.item()-1) < 1e-4
# This is a mechanics check on five points, not held-out performance evidence.
