Dataset、DataLoader 与训练循环：从批训练到独立推理

from torch.utils.data import Dataset


class PairDataset(Dataset):
    def __init__(self, features, labels):
        self.features = features
        self.labels = labels

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, index):
        return self.features[index], self.labels[index]

from pathlib import Path

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset, random_split

torch.manual_seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

features = torch.randn(600, 4)
noise = torch.randn(600) * 0.35
labels = (
    features[:, 0] - 0.8 * features[:, 1]
    + 0.5 * features[:, 2] ** 2 + noise > 0.4
).long()
dataset = TensorDataset(features, labels)
train_data, valid_data = random_split(
    dataset, [480, 120], generator=torch.Generator().manual_seed(42)
)
train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
valid_loader = DataLoader(valid_data, batch_size=64, shuffle=False)


class RiskNet(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 2)
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.layers(inputs)


def train_epoch(model, loader, loss_fn, optimizer) -> float:
    model.train()
    total_loss = 0.0
    for batch_features, batch_labels in loader:
        batch_features = batch_features.to(device)
        batch_labels = batch_labels.to(device)
        optimizer.zero_grad()
        loss = loss_fn(model(batch_features), batch_labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * len(batch_labels)
    return total_loss / len(loader.dataset)


def evaluate(model, loader, loss_fn) -> tuple[float, float]:
    model.eval()
    total_loss, correct = 0.0, 0
    with torch.inference_mode():
        for batch_features, batch_labels in loader:
            batch_features = batch_features.to(device)
            batch_labels = batch_labels.to(device)
            logits = model(batch_features)
            total_loss += loss_fn(logits, batch_labels).item() * len(batch_labels)
            correct += (logits.argmax(dim=1) == batch_labels).sum().item()
    return total_loss / len(loader.dataset), correct / len(loader.dataset)


model = RiskNet().to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
checkpoint = Path("model_output/risk_net.pt")
checkpoint.parent.mkdir(parents=True, exist_ok=True)
best_loss = float("inf")

for epoch in range(1, 31):
    train_loss = train_epoch(model, train_loader, loss_fn, optimizer)
    valid_loss, valid_accuracy = evaluate(model, valid_loader, loss_fn)
    if valid_loss < best_loss:
        best_loss = valid_loss
        torch.save(model.state_dict(), checkpoint)
    if epoch in (1, 10, 20, 30):
        print(epoch, round(train_loss, 4), round(valid_loss, 4), round(valid_accuracy, 3))

loaded = RiskNet().to(device)
loaded.load_state_dict(torch.load(checkpoint, map_location=device, weights_only=True))
loaded.eval()
with torch.inference_mode():
    sample_probability = loaded(features[:1].to(device)).softmax(dim=1)
print("probability:", sample_probability.cpu().numpy().round(3))