Pytorch函数集合

这是一个笔记本，记录在pytorch种遇到的一些小函数的使用说明，不断更新中。

gt、ge、eq、lt、le系列函数：

torch.gt(input,value)
Returns: bool tensor,shape as input

即greaterthan、greateerqual、equal、lessthan、lessequal

用于比较tensor与指定值的大小，输出同shape的bool型变量

torch.long、torch.uint8类型用作索引的区别:

t = torch.randn(4,2)
mask = torch.ones(4,dtype=torch.uint8)
mask[2] = 0
print(t[mask,:])

uint8可以理解为bool型索引，此时需要与tensor的维度保持一致，当对应位置的元素为1或True时则从t中输出该元素

t = torch.randn(4,2)
mask = torch.ones(4,dtype=torch.long)
mask[2] = 0
print(t[mask,:])

long可以理解为要取出的元素的索引，无需与tensor维度一致，上述代码从t中依次取出[1,1,0,1]位置的元素

Tensor.repeat,torch.meshgrid:

t = torch.rand(1,3)
t_r = t.repeat(5,1)
print(t_r.shape)

输出为torch.size([5,3])，将tensor在对应位置上重复n次。

例如：想要得到5x3大小的特征图上每个点的xy坐标，则x为,y为

对于每列上的x坐标都是相同的，所以生成一列x坐标，然后重复3次

那么首先生成一列x坐标:cx = torch.arange(5).view(5,1)

然后在列方向重复3次变成3列:cx = cx.repeat(1,3)

最终得到了(5,3)的cx坐标

y坐标类似，就是torch.arange(3).view(1,3).repeat(5,1)，得到(5,3)的cy坐标

最后用torch.stack([a,b],axis=2)即可得到(5,3,2)的cxy坐标

注意：把矩阵看作图像时索引的xy是反的。

该图为(3,3)的矩阵A,取红色方框时为A[3,2],按照图像的xy坐标则为[2,3].(设索引从1开始）

另一种方式就是torch.meshgrid函数：

torch.meshgrid函数是按照矩阵的方式生成网格坐标的，例如特征图为5x3大小，即宽W=3，高H=5

x,y = torch.meshgrid(torch.arange(H),torch.arange(W))

就生成(5x3)的网格矩阵了。

np.meshgrid则是按照图像的xy方式生成网格的。torch更符合直觉

nn.BCELoss,nn.BCEWithLogitsLoss,nn.CrossEntropyLoss

nn.CrossEntropyLoss的输入为：preds->(batch,n_class,feat_size),labels(batch,feat_size)

函数会先对preds作softmax，然后再进行交叉熵损失函数的计算

nn.BCELoss和WithLogitsLoss输入的preds和labels维度相同，都是->(batch,C,feat_size),

函数会在C维度上进行进项逐项的交叉熵损失计算

WithLogitsLoss会先进行sigmoid再计算交叉熵。

nn.BCEWithLogitsLoss(pos_weight=torch.Tensor())

pos_weight参数用来调整不同类损失所占权重。

例如：3类的分类任务，则网络输出为(n,3),设置reduciton='none'，则损失输出为(n,3)

BCEcls_w = nn.BCEWithLogitsLoss(pos_weight = torch.Tensor([0.1,0.2,0.3]),reduction='none')
BCEcls = nn.BCEWithLogitsLoss(pos_weight = torch.Tensor([1]),reduction='none')
label = torch.Tensor([[1,0,0],[0,0,1]])
pred = torch.Tensor([[1,0.1,0],[0,1,3]])
BCEcls(pred,label)
#tensor([[0.3133, 0.7444, 0.6931],
        [0.6931, 1.3133, 0.0486]])
BCEcls_w(pred,label)
#tensor([[0.0313, 0.7444, 0.6931],
        [0.6931, 1.3133, 0.0146]])

 并不是对3个数字都乘weight，样本1预测结果为0类，只对0类的预测乘weight，其他保持不变。

类似的，样本2预测结果为2类，则只对0.0486乘权重0.3。

pytorch NMS函数

def box_iou(box1, box2):
    # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
    """
    Return intersection-over-union (Jaccard index) of boxes.
    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
    Arguments:
        box1 (Tensor[N, 4])
        box2 (Tensor[M, 4])
    Returns:
        iou (Tensor[N, M]): the NxM matrix containing the pairwise
            IoU values for every element in boxes1 and boxes2
    """

    def box_area(box):
        # box = 4xn
        return (box[2] - box[0]) * (box[3] - box[1])

    area1 = box_area(box1.T)
    area2 = box_area(box2.T)

    # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
    inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
    return inter / (area1[:, None] + area2 - inter)  # iou = inter / (area1 + area2 - inter)

Params:
待NMS的检测框sets,对应的得分scores
# 1表示被抑制，0表示不抑制
suppress = torch.zeros(sets.size(0),dtype=torch.bool)
_,idx = scores.sort(descending=True)
sets_sorted = sets[idx]
ious = box_iou(sets_sorted,sets_sorted)
for k in range(ious.size(0)):
    if suppress[k] == 1:
        continue
    suppress = torch.max(suppress, (ious[k]>threshod))
    #上面一步会把当前框也抑制掉，所以强制改为0
    suppress[k] = 0
#1-suppress就是NMS后剩余框的位置，即[True,False...]
idx[~suppress]  #即可取出保留框的索引