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]