通常,由於類別不均衡,需要使用weighted cross entropy loss平衡。
def inverse_freq(label):
"""
輸入label [N,1,H,W],1是channel數目
"""
den = label.sum() # 0
_,_,h,w= label.shape
num = h*w
alpha = den/num # 0
return torch.tensor([alpha, 1-alpha]).cuda()
# train
...
loss1 = F.cross_entropy(out1, label.squeeze(1).long(), weight=inverse_freq(label))
代碼比較簡單,寫在博客上保存。