pytorch中交叉熵函数torch.nn.CrossEntropyLoss()怎么加入权重

关于交叉熵函数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)

方法二:

### 我的类别数为12
criterion = nn.CrossEntropyLoss(weight=torch.FloatTensor([0.1,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0]) ,
                                size_average=True)
### 要想使用GPU,此处必须使用cuda()
criterion.cuda()   
### y为预测结果, tags为原标签
loss = criterion(y, tags)

参考自:
https://blog.csdn.net/lidichengfo0412/article/details/97786439

猜你喜欢

转载自blog.csdn.net/yuekangwei/article/details/111491256