pytorch使用小技巧

给定预测pred维度[n, c]，target[n]，根据真值类别来得到预测值里对应的概率

pt = pred.gather(1, target)  # 1表示在pred指定维度根据target值获取相关类别预测概率

按照指定维度分割x，[3, 2]表示分成两个张量在指定维度的长度，一个张量A[1, 3, 352, 320],一个张量B[1, 2, 352, 320]

A, B = torch.split(X, [3,2], dim=1)  # X的维度[1, 5, 352, 320]

第二种写法是不对的，会导致self.weights没有得到梯度更新，因为.cuda作用于tensor会把其移动到GPU上，但是作用于Parameter时只会在GPU上拷贝一个副本，在梯度更新时，并不会影响到原始的Parameter参数。 

self.weights = nn.Parameter(weights.float().cuda())
self.weights = nn.Parameter(weights.float()).cuda()