【Pytorch学习笔记】模型模块05——Module常用函数

Module常用函数

设置训练和评估模式

**作用：**在PyTorch中，模型有训练(training)和评估(evaluation)两种模式，它们会影响某些层的行为。

主要影响的层：

• Dropout层：训练时随机丢弃神经元，评估时保持全部神经元
• BatchNorm层：训练时计算并更新统计量，评估时使用固定统计量
• LayerNorm层：行为在两种模式下基本一致

2. 设置方法

# 设置训练模式
model.train()

# 设置评估模式
model.eval()

3. 实际应用示例

# 训练循环
model.train()
for epoch in range(num_epochs):
    for batch_x, batch_y in train_loader:
        optimizer.zero_grad()
        output = model(batch_x)
        loss = criterion(output, batch_y)
        loss.backward()
        optimizer.step()

# 评估循环
model.eval()
with torch.no_grad():  # 禁用梯度计算
    for batch_x, batch_y in val_loader:
        output = model(batch_x)
        # 计算验证指标

4. 注意事项

• 切换模式会影响所有子模块
• 评估时建议使用 torch.no_grad() 来提高效率
• 某些自定义层可能需要重写 train() 和 eval() 方法
• 在使用预训练模型时，注意保持正确的模式设置

**常见错误：**忘记在训练和评估阶段之间切换模式，可能导致性能下降或结果不准确。

模型设备设置

1. 设备类型

• **CPU：**默认设备，适用于小型模型和数据集
• **GPU：**加速计算，适用于大型模型和数据集训练

2. 设置方法

# 检查GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 将模型移动到指定设备
model = model.to(device)

# 将数据移动到指定设备
inputs = inputs.to(device)
labels = labels.to(device)

3. 实际应用示例

# 完整训练示例
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MyModel().to(device)

for batch_x, batch_y in train_loader:
    # 将数据移到相应设备
    batch_x = batch_x.to(device)
    batch_y = batch_y.to(device)
    
    outputs = model(batch_x)
    loss = criterion(outputs, batch_y)

4. 注意事项

• 模型和数据必须在同一设备上才能进行计算
• 使用多GPU时可以使用 DataParallel 或 DistributedDataParallel
• 某些操作（如打印、可视化）可能需要将数据移回CPU
• 使用 .to(device) 方法不会改变原始数据，而是返回新的副本

**最佳实践：**在训练开始时就确定设备，并始终保持数据和模型在同一设备上。

获取和加载模型参数

1. state_dict() 方法

**作用：**获取模型的所有参数（权重和偏置）及其当前值。返回一个有序字典，包含每一层的参数张量。

# 获取模型参数
model_state = model.state_dict()

# 保存模型参数
torch.save(model_state, 'model_weights.pth')

2. load_state_dict() 方法

**作用：**将保存的参数加载到模型中，用于模型的权重恢复或迁移学习。

# 加载保存的参数
model_state = torch.load('model_weights.pth')
model.load_state_dict(model_state)

3. 实际应用示例

# 完整的保存和加载示例
# 保存模型
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, 'checkpoint.pth')

# 加载模型
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

4. 注意事项

• 确保加载的状态字典与模型结构匹配
• 可以使用 strict=False 参数在部分加载时忽略不匹配的参数
• 在多GPU训练时注意处理 DataParallel 包装的模型状态字典
• 加载预训练模型时可能需要调整键名以匹配目标模型结构

**最佳实践：**定期保存模型检查点，包含额外信息（如训练轮次、优化器状态等）以便恢复训练。

管理模型的modules、parameters、sub_module

1. parameters() 和 named_parameters()

**作用：**获取模型中所有可训练的参数。

# parameters() 返回参数迭代器
for param in model.parameters():
    print(param.shape)

# named_parameters() 返回参数名称和值的迭代器
for name, param in model.named_parameters():
    print(f"Layer: {name} | Size: {param.shape}")

2. children() 和 named_children()

**作用：**获取模型的直接子模块（不包括子模块的子模块）。

# children() 返回直接子模块的迭代器
for child in model.children():
    print(child)

# named_children() 返回子模块名称和模块的迭代器
for name, child in model.named_children():
    print(f"Child name: {name}\\nChild module: {child}\\n")

3. modules() 和 named_modules()

**作用：**获取模型的所有模块（包括子模块的子模块）。

# modules() 返回所有模块的迭代器
for module in model.modules():
    print(module)

# named_modules() 返回模块名称和模块的迭代器
for name, module in model.named_modules():
    print(f"Module name: {name}\\nModule: {module}\\n")

4. get_parameter() 和 get_submodule()

**作用：**通过名称获取特定的参数或子模块。

# 获取特定参数
param = model.get_parameter('conv1.weight')

# 获取特定子模块
submodule = model.get_submodule('encoder.layer1')

5. add_module()

**作用：**向模型动态添加新的子模块。

# 添加新模块
model.add_module('new_layer', nn.Linear(100, 10))

# 实际应用示例
class DynamicModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.add_module('conv1', nn.Conv2d(1, 20, 5))
        self.add_module('relu1', nn.ReLU())

6. 注意事项

• parameters() 和 named_parameters() 只返回需要梯度的参数
• children() 和 modules() 的区别在于遍历深度
• 使用 get_parameter() 和 get_submodule() 时要确保名称正确
• add_module() 添加的模块名称不能重复

**最佳实践：**在调试或修改模型结构时，使用这些方法可以方便地检查和操作模型的各个组件。

设置模型的参数精度

1. 精度类型概述

• float32 (float): 默认精度类型，32位浮点数，适用于大多数场景
• float16 (half): 16位浮点数，可减少内存使用和计算时间
• bfloat16: Brain Floating Point，16位浮点数的变体，具有更好的数值稳定性
• float64 (double): 64位浮点数，提供更高的精度，但消耗更多资源

2. 设置方法

# float32
model = model.float()

# float16
model = model.half()

# bfloat16
model = model.bfloat16()

# float64
model = model.double()

3. 实际应用示例

# 使用混合精度训练
from torch.cuda.amp import autocast, GradScaler

# 初始化梯度缩放器
scaler = GradScaler()

# 训练循环
for batch_x, batch_y in train_loader:
    # 使用自动混合精度
    with autocast():
        output = model(batch_x)
        loss = criterion(output, batch_y)
    
    # 使用梯度缩放器处理反向传播
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

4. 注意事项

• 降低精度可能会影响模型性能和训练稳定性
• 某些运算可能不支持特定精度类型
• 使用混合精度训练可以平衡计算效率和数值稳定性
• 在使用低精度时，需要特别注意数值溢出问题

**最佳实践：**对于大型模型，建议使用混合精度训练来平衡训练速度和模型性能。在推理阶段，可以根据具体需求选择合适的精度类型。

对子模块执行特定功能

1. apply() 方法

**作用：**将一个函数递归地应用到模型的每个子模块上。

# 权重初始化示例
def init_weights(m):
    if isinstance(m, nn.Linear):
        torch.nn.init.xavier_uniform_(m.weight)
        m.bias.data.fill_(0.01)

# 应用到整个模型
model.apply(init_weights)

2. zero_grad() 方法

**作用：**将模型中所有可训练参数的梯度设置为零。通常在每次反向传播前调用。

在PyTorch中，梯度会默认累积。如果不清零，每次反向传播计算的梯度会与之前的梯度相加。这可能导致：

• 梯度累积导致更新步长过大，影响模型训练稳定性
• 上一次迭代的梯度会影响当前批次的优化方向，使得模型收敛方向不准确

因此，在每次反向传播前调用zero_grad()来清零梯度是非常重要的。这样可以确保每个批次的梯度计算都是独立的，不受之前迭代的影响。

# 训练循环中使用
for batch_x, batch_y in train_loader:
    optimizer.zero_grad()  # 清零梯度
    outputs = model(batch_x)
    loss = criterion(outputs, batch_y)
    loss.backward()
    optimizer.step()

3. requires_grad_() 方法

**作用：**设置模型参数是否需要计算梯度。用于微调或冻结某些层。

# 冻结所有参数
for param in model.parameters():
    param.requires_grad_(False)

# 只训练最后几层
for param in model.classifier.parameters():
    param.requires_grad_(True)

4. 注意事项

• apply() 方法会递归访问所有子模块，执行效率需要考虑
• 在使用 zero_grad() 时，确保在正确的时机调用以避免梯度累积
• 使用 requires_grad_() 时要注意梯度流的传播路径

**最佳实践：**根据具体任务需求选择合适的方法，并在正确的时机调用这些方法以确保模型正常工作。