PyTorch Demo

线性模型

import torch
from torch import nn


def training_loop(p_epochs, p_optimizer, p_model, p_loss_fn, p_train_u, p_val_u, p_train_c, p_val_c):
    for epoch in range(p_epochs):
        train_loss = p_loss_fn(p_model(p_train_u), p_train_c)
        val_loss = p_loss_fn(p_model(p_val_u), p_val_c)
        p_optimizer.zero_grad()
        train_loss.backward()
        p_optimizer.step()
        if epoch % 1000 == 0:
            print(f"epoch {epoch + 1}, train_loss {train_loss.item():.4f}, val_loss {val_loss:.4f}")


t = torch.rand(100).unsqueeze_(1)
c = t * 2 + 1
shuffled_indices = torch.randperm(t.shape[0])
train_indices, val_indices = shuffled_indices[:-int(0.2 * t.shape[0])], shuffled_indices[-int(0.2 * t.shape[0]):]
linear_model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(linear_model.parameters(), lr=1e-2)
training_loop(10000, optimizer, linear_model, nn.MSELoss(), t[train_indices], t[val_indices], c[train_indices],
              c[val_indices])
print(linear_model.weight)
print(linear_model.bias)