pytorch的交叉熵loss分权重nn.CrossEntropyLoss——训练集各类数量不均衡的解决——样本数量不平衡

仅作为记录，大佬请跳过。

感谢大佬师兄的提示，对网络的loss对各类设置权重（只用一行代码），解决训练集各类数量不均衡的问题（0类：49w，1类：7w；4类：155w）；（测试集的数据各类可以不均衡）

代码

# ********************************************************************************************************************************
# criterion = nn.CrossEntropyLoss()
import glob
imglist_dir_0=data_dir+'/train/0/';imglist_dir_1=data_dir+'/train/1/';imglist_dir_2=data_dir+'/train/2/';imglist_dir_3=data_dir+'/train/3/';imglist_dir_4=data_dir+'/train/4/'
num_rev_0,num_rev_1,num_rev_2,num_rev_3,num_rev_4=1/(len(glob.glob(imglist_dir_0+'*.png'))),1/(len(glob.glob(imglist_dir_1+'*.png'))),1/(len(glob.glob(imglist_dir_2+'*.png'))),1/(len(glob.glob(imglist_dir_3+'*.png'))),1/(len(glob.glob(imglist_dir_4+'*.png')))
criterion = nn.CrossEntropyLoss(weight=torch.from_numpy(np.array([num_rev_0,num_rev_1,num_rev_2,num_rev_3,num_rev_4])).float() ,
                                size_average=True)
criterion.to(device)
# **************************************************** END ****************************************************************************

插入网络中的展示

注：要加入criterion.to(device)

参考

感谢大佬博主文章传送门

传送门

————————————————————

简示例代码

import glob
data_dir = '/Data4T/dataset_her2/dataset_her2_server_five-class_new_v2'
imglist_dir_0=data_dir+'/train/0/';imglist_dir_1=data_dir+'/train/1/';imglist_dir_2=data_dir+'/train/2/';imglist_dir_3=data_dir+'/train/3/';imglist_dir_4=data_dir+'/train/4/'
num_rev_0,num_rev_1,num_rev_2,num_rev_3,num_rev_4=1/(len(glob.glob(imglist_dir_0+'*.png'))),1/(len(glob.glob(imglist_dir_1+'*.png'))),1/(len(glob.glob(imglist_dir_2+'*.png'))),1/(len(glob.glob(imglist_dir_3+'*.png'))),1/(len(glob.glob(imglist_dir_4+'*.png')))
import torch
import torch.nn as nn
import numpy as np
criterion = nn.CrossEntropyLoss(weight=torch.from_numpy(np.array([num_rev_0,num_rev_1,num_rev_2,num_rev_3,num_rev_4])).float() ,
                                size_average=True)
#criterion.to(device)
print(num_rev_0)
print(num_rev_1)
print(num_rev_2)
print(num_rev_3)
print(num_rev_4)

简示例展示