Pytorch中模型常用片段汇总

文章目录

• 前言
• 1、统计模型的可学习参数的数量
• 2、模型不同位置使用不同学习率
• 3、目标检测中box的变换
• 4、张量提取 *** 作
• 5、绘制边框和坐标
• 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

 比如conditional detr的参数数目为43,196,601即可学习参数的数量为43M。

---

2、模型不同位置使用不同学习率

 比如给backbone和model其余部分采用不同的学习率进行更新：

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 
backbone.0.body.layer4.0.conv1.weight 
'''
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()

 上述部分中单独指定了backbone的学习率为1e-5，其余模块学习率为1e-4，且权重衰减为1e-4；且训练经过40轮epochs后更新一次学习率。

另外，还设定了一个梯度裁剪参数，若模型的的参数梯度超过max_norm则执行裁剪。

3、目标检测中box的变换

 在目标检测算法中，经常涉及boxes由[cx,cy,w,h] <–> [x1,y1,x2,w2]以及rescale的 *** 作变换。

在此提供两个函数：
 参考地址：Facebook_detr
 code:

def box_cxcywh_to_xyxy(x):
    x_c, y_c, w, h = x.unbind(1)   # [N]
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
         (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=1)   # [N] --> [N,1] --> [N,4]

def rescale_bboxes(out_bbox, size):
    img_w, img_h = size
    b = box_cxcywh_to_xyxy(out_bbox)
    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
    return b

# boxes： [cx,cy,w,h]
boxes =torch.tensor([[1,2,3,4],[5,6,7,8]],dtype=torch.float32)
# rescale
res = rescale_bboxes(boxes,(1,1))
print(res)

4、张量提取 *** 作

 常用来提取置信度>某个阈值的 *** 作，语言描述不出来，看代码吧：

# 若batch == 1
boxes = torch.randn(1,5,4) # [b=1,num_boxes,4]
keep = torch.tensor([0,0,0,1,1]).to(torch.bool) # [num_boxes]的bool型张量
res = boxes[0,keep]        # [2,4]，将keep中TRUE的boxes提取出来，注意0不能舍去
print(res)
print(res.shape)
# 若batch>1，常用来计算loss，因为可以消除掉batch这个维度
boxes = torch.randn(2,5,4)  # [batch, num_boxes, 4]
keep = torch.tensor([[0,0,0,1,1],[1,1,1,0,0]]).to(torch.bool) # [batch,num_boxes]
res = boxes[keep]           # [num_truth,4]
print(res)
print(res.shape)

 简单理解，若想提取布尔型张量keep中的TRUE所对应的元素，则keep的shape必须和待提取张量的前两维度相同：4维张量则3维的keep；3维张量则2维的keep。

5、绘制边框和坐标

 同样，在facebook_detr搬运过来的。

def plot_results(pil_img, prob, boxes):
	'''
	pil_img: Image.open()对象
	prob:[num_TP,num_classes]的Tensor
	boxes:[num_TP,4]
	'''
    plt.figure(figsize=(16,10))
    plt.imshow(pil_img)
    ax = plt.gca()
    for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), COLORS * 100):
        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
                                   fill=False, color=c, linewidth=3))
        cl = p.argmax()
        text = f'{CLASSES[cl]}: {p[cl]:0.2f}'
        ax.text(xmin, ymin, text, fontsize=15,
                bbox=dict(facecolor='yellow', alpha=0.5))
    plt.axis('off')
    plt.show()

N、分布式相关