深度学习中模型打印每个迭代的损失值和梯度范数

前言

为了打印每个迭代的损失值和梯度范数，需要在训练循环中添加相应的代码。
这里是一个修改后的训练函数，它包含了打印损失值和梯度范数的逻辑：

代码

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

def train(model, train_loader, test_loader, num_epochs, loss_function, optimizer, args):
    # 初始化训练和测试loss的列表
    train_losses = []
    
    # 加载预训练模型（仅一次）并冻结其参数
    pretrained_model = torch.load("./model_saved/dnn_forward.pt").to(args.device)
    pretrained_model.train()
    
    # 训练循环
    model.train()
    for epochs in range(num_epochs):
        epoch_loss = 0.0
        for i, (x_train, y_train) in enumerate(train_loader):
            x_train, y_train = x_train.to(args.device), y_train.to(args.device)
            optimizer.zero_grad()  # 清除之前的梯度
            
            y_pred = model(x_train)  # 当前模型的前向传播
            with torch.no_grad():
                dnn_p = pretrained_model(y_pred)  # 通过预训练模型
            
            loss = loss_function(dnn_p, x_train)  # 计算损失
            loss.backward()  # 反向传播计算梯度
            
            # 打印损失值和梯度范数
            print(f"Epoch [{epochs+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], "
                  f"Loss: {loss.item():.4f}, "
                  f"Grad Norm: {calculate_grad_norm(model.parameters()):.4f}")
            
            optimizer.step()  # 更新模型参数
            epoch_loss += loss.item() * x_train.size(0)  # 累计损失
        
        # 计算并保存平均损失
        epoch_loss /= len(train_loader.dataset)
        train_losses.append(epoch_loss)
        print(f"Epoch [{epochs+1}/{num_epochs}] average Loss: {epoch_loss:.4f}")
    
    # 如果有测试集，可以在这里添加测试代码
    # ...
    
    return train_losses  # 返回训练损失列表

# 辅助函数：计算梯度范数
def calculate_grad_norm(parameters):
    total_norm = 0.0
    for p in parameters:
        if p.grad is not None:
            param_norm = p.grad.data.norm(2)
            total_norm += param_norm.item() ** 2
    total_norm = total_norm ** (1. / 2)
    return total_norm

• 在这个修改后的版本中，calculate_grad_norm 函数被用来计算模型中所有参数的梯度范数。
• 这个范数是通过将每个参数的梯度范数平方相加，然后取平方根得到的。
• 在训练循环的每次迭代中，都会打印出当前的迭代步数、损失值和梯度范数。

小结

• 请注意，打印太多的信息可能会使训练过程变慢，特别是当数据集很大或迭代次数很多时。
• 因此，可能想要根据实际需要调整打印的频率，例如每隔一定的迭代步数打印一次信息。