不知道如何记录模型训练的log？一文带你入门

loguru对AI模型指标的记录

代码如下

import torch
from loguru import logger
from torch.utils.data import DataLoader
from torchvision import datasets
import torchvision.transforms as transforms
from torchvision.models import resnet50
from torch.optim import SGD
from torch.nn import CrossEntropyLoss
def train_model(model, data_loader, optimizer, criterion, num_epochs):
    for epoch in range(num_epochs):
        running_loss = 0.0
        correct_predictions = 0
        total_predictions = 0

        for inputs, labels in data_loader:
            # 梯度清零
            optimizer.zero_grad()

            # 前向传播
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            # 反向传播和优化
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

            # 计算准确率
            _, predicted = torch.max(outputs.data, 1)
            total_predictions += labels.size(0)
            correct_predictions += (predicted == labels).sum().item()

        # 计算平均损失和准确率
        average_loss = running_loss / len(data_loader)
        accuracy = correct_predictions / total_predictions

        # 使用loguru记录每个epoch的指标
        logger.info(f"Epoch {epoch + 1}/{num_epochs}:")
        logger.info(f"	Average Loss: {average_loss}")
        logger.info(f"	Accuracy: {accuracy}")

if __name__ == "__main__":
    # 这里假设已经有定义好的模型、数据加载器、优化器和损失函数
    model =resnet50()
    data_loader = DataLoader(datasets.CIFAR10(root='./data', train=True, download=True, transform=transforms.ToTensor()), batch_size=16, shuffle=True)
    optimizer =SGD(model.parameters(), lr=0.001)
    criterion =CrossEntropyLoss()
    num_epochs = 10

    train_model(model, data_loader, optimizer, criterion, num_epochs)

结果展示

log记录可视化

• tensorboard
• wandb