关于交叉熵函数torch.nn.CrossEntropyLoss的基本用法可参考https://blog.csdn.net/zziahgf/article/details/80196376。
那么如何在交叉熵函数中添加自定义的各类别的权重呢?定义参数weight即可,其中参数的内容要为tensor类型,而且要把dtype类型定义为float。用法为:
### 我的类别数为12
criterion = nn.CrossEntropyLoss(weight=torch.from_numpy(np.array([0.1,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])).float() ,
size_average=True)
### 要想使用GPU,此处必须使用cuda()
criterion.cuda()
### y为预测结果, tags为原标签
loss = criterion(y, tags)