PyTorch 训练技巧与模型微调指南

断点恢复

# ============================ step 2/5 模型 ============================
net = LeNet(classes=2)
net.initialize_weights()
# ============================ step 3/5 损失函数 ============================
criterion = nn.CrossEntropyLoss()
# 选择损失函数
# ============================ step 4/5 优化器 ============================
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9)
# 选择优化器
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=6, gamma=0.1)
# 设置学习率下降策略
# ============================ step 5+/5 断点恢复 ============================
path_checkpoint = "./checkpoint_4_epoch.pkl"
checkpoint = torch.load(path_checkpoint)
net.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch']
scheduler.last_epoch = start_epoch

2. 模型微调

模型微调包括 3 个步骤：1.构建模型；2.加载参数；3.修改全连接 fc 层

# 3/3 替换 fc 层
num_ftrs = resnet18_ft.fc.in_features
resnet18_ft.fc = nn.Linear(num_ftrs, classes)

模型微调方法包括两种：

1.固定预训练的参数（requires_grad=False, lr=0）

2.Features Extractor 较小学习率（params_group）

2.1 固定预训练的参数（requires_grad=False,lr=0）

# ============================ step 2/5 模型 ============================
# 1/3 构建模型
resnet18_ft = models.resnet18()
# 2/3 加载参数
# flag = 0
flag = 1
if flag:
    path_pretrained_model = "pretrained_model.pth"  # 通用路径
    state_dict_load = torch.load(path_pretrained_model)
    resnet18_ft.load_state_dict(state_dict_load)
# 法 1 : 冻结卷积层
flag_m1 = 0
# flag_m1 = 1
if flag_m1:
    for param in resnet18_ft.parameters():
        param.requires_grad = False
    print("conv1.weights[0, 0, ...]:\n {}".format(resnet18_ft.conv1.weight[0, 0, ...]))
# 3/3 替换 fc 层
num_ftrs = resnet18_ft.fc.in_features
resnet18_ft.fc = nn.Linear(num_ftrs, classes)
resnet18_ft.to(device)
# ============================ step 3/5 损失函数 ============================
criterion = nn.CrossEntropyLoss()
# 选择损失函数
# ============================ step 4/5 优化器 ============================
optimizer = optim.SGD(resnet18_ft.parameters(), lr=LR, momentum=0.9)
# 选择优化器
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=lr_decay_step, gamma=0.1)
# 设置学习率下降策略

2.2 Features Extractor 较小学习率（params_group）

# ============================ step 2/5 模型 ============================
# 1/3 构建模型
resnet18_ft = models.resnet18()
# 2/3 加载参数
# flag = 0
flag = 1
if flag:
    path_pretrained_model = "pretrained_model.pth"  # 通用路径
    state_dict_load = torch.load(path_pretrained_model)
    resnet18_ft.load_state_dict(state_dict_load)
# 3/3 替换 fc 层
num_ftrs = resnet18_ft.fc.in_features
resnet18_ft.fc = nn.Linear(num_ftrs, classes)
resnet18_ft.to(device)
# ============================ step 3/5 损失函数 ============================
criterion = nn.CrossEntropyLoss()
# 选择损失函数
# ============================ step 4/5 优化器 ============================
# 法 2 : conv 小学习率
fc_params_id = list(map(id, resnet18_ft.fc.parameters()))
base_params = filter(lambda p: id(p) not in fc_params_id, resnet18_ft.parameters())
optimizer = optim.SGD([{'params': base_params, 'lr': LR*0}, {'params': resnet18_ft.fc.parameters(), 'lr': LR}], momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=lr_decay_step, gamma=0.1)
# 设置学习率下降策略

使用第二种方法更加灵活。

3. GPU 的使用

3.1 CPU 与 GPU 及数据在 CPU 与 GPU 间的迁移

pytorch 有两种数据类型：tensor 和 module。data=tensor+module.

从 CPU 到 GPU:data.to("cuda")

从 GPU 到 CPU:data.to("cpu")

to() 函数，数据类型转换或者设备转换。执行 to 函数是张量不执行 replace，module 执行 replace。

3.2 GPU 并行运算

4.pytorch 常见问题