Pytorch-Tricks

class BadEncoder(nn.Module):
    def forward(self, x):
        if x.sum() > 0:          # 数据依赖的条件
            x = self.heavy_conv(x)
        else:
            x = self.light_conv(x)
        return x

class BadRNN(nn.Module):
    def forward(self, x):
        for i in range(x.size(1)):   # 依赖输入序列长度
            x = self.rnn_cell(x[:, i, :], x)
        return x

import torch
import torch.nn as nn
import torch.optim as optim
from torch.profiler import profile, record_function, ProfilerActivity, tensorboard_trace_handler

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1000, 1000)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(1000, 10)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleModel().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=tensorboard_trace_handler('./log/profiler_result'),
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    
    for i in range(10):
        inputs = torch.randn(32, 1000).to(device)
        labels = torch.randint(0, 10, (32,)).to(device)

        with record_function("model_forward"):
            outputs = model(inputs)
            
        with record_function("loss_backward"):
            loss = criterion(outputs, labels)
            optimizer.zero_grad()
            loss.backward()

        with record_function("optimizer_step"):
            optimizer.step()

        prof.step()

# 按 GPU 总耗时排序
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
# 按 CPU 总耗时排序
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

def forward(self, x, epoch):
    x = layer1(x)
    if epoch < 5:
        # 冻结部分参数
        with torch.no_grad():
            x = self.layer2(x)
    else:
        # 正常训练
        x = self.layer2(x)
    return x

RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by (1) passing the keyword argument `find_unused_parameters=True` to `torch.nn.parallel.DistributedDataParallel`; (2) making sure all `forward` function outputs participate in calculating loss. If you already have done the above two steps, then the distributed data parallel module wasn't able to locate the output tensors in the return value of your module's `forward` function. Please include the loss function and the structure of the return value of `forward` of your module when reporting this issue (e.g. list, dict, iterable).

model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters=True)

# forward 函数中确保所有参数都参与计算
def forward(self, x, epoch):
    x = self.layer1(x)
    x = self.layer2(x)       # 始终正常前向，不断开梯度
    x = self.layer3(x)
    return x

# 冻结 layer2 的参数
optimizer = torch.optim.Adam([
    {'params': model.module.layer2.parameters(), 'lr': 0.0},   # 前5个epoch不更新
    {'params': model.module.layer1.parameters()},
    {'params': model.module.layer3.parameters()},
], lr=1e-3)

# epoch 5 后更新 layer2 的参数
for param_group in optimizer.param_groups:
    if param_group['params'][0] is model.module.layer2.parameters().__next__():  # 简单判断
        param_group['lr'] = 1e-3

# 冻结参数
for param in model.layer2.parameters():
    param.requires_grad = False
# 训练前设置 DDP
model = torch.nn.parallel.DistributedDataParallel(model)
# 5 epoch 后解冻参数
for param in model.layer2.parameters():
    param.requires_grad = True
# 重新初始化 DDP
model = torch.nn.parallel.DistributedDataParallel(model)