Pytorch中模型常用片段彙總

文章目錄

• 前言
• 1、統計模型的可學習參數的數量
• 2、模型不同位置使用不同學習率
• N、分布式相關

前言

 本文主要整理一些pytorch關于搭模組化型過程中常用的代碼片段。

1、統計模型的可學習參數的數量

 這個名額是我在conditional detr論文中看見的，即比較的是detr和conditional detr兩個方法參數的數量。注意，此處統計的是數量，并不是參數所占有的記憶體大小。貼上公開源碼：

n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
print('number of params:', n_parameters)  # 43,196,001
           

2、模型不同位置使用不同學習率

param_dicts = [
	# 除backbone其餘子產品的可學習參數，n是名字，屬于字元串對象
    {"params": [p for n, p in model_without_ddp.named_parameters() if "backbone" not in n and p.requires_grad]},
    # backbone單獨指定的學習率
    {
        "params": [p for n, p in model_without_ddp.named_parameters() if "backbone" in n and p.requires_grad],
        "lr": 1e-5,
    },
]
'''
n:
backbone.0.body.layer3.5.conv3.weight <class 'str'>
backbone.0.body.layer4.0.conv1.weight <class 'str'>
'''
optimizer = torch.optim.AdamW(param_dicts, lr=1e-4, weight_decay=1e-4)
# 40輪epochs後更新一次學習率。
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 40)

# 在每一個iter後，更新一次優化器
for cur_epoch in epochs:
	for data in cur_epoch:
	    optimizer.zero_grad()
	    losses.backward()
	    # max_norm = 0.1是個超參。
	    # 執行梯度裁剪
	    if max_norm > 0:
	        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
	    optimizer.step()
	# 更新一輪epoch後，在調用lr_scheduler。
	lr_scheduler.step()
           

N、分布式相關