stack expects each tensor to be equal size, but got XXX at entry X and YYY at entry Y的debug方法

这一条bug的出现来源于CV的数据处理阶段,同一batch内图像或标签的shape不一致,导致dataloader阶段会出现报错

推荐解决方法:通过复写dataloader终端 collate_fn参数对输入的数据进行调整,使之能具有相同的shape,collate_fn的作用是将输入的数据通过特定的连接整合方法使其成为batch的格式,重写collate_fn相当于对batch的数据格式进行自定义

def collate_fn(self,data):
    imgs_list,boxes_list,classes_list=zip(*data)
    assert len(imgs_list)==len(boxes_list)==len(classes_list)
    batch_size=len(boxes_list)
    pad_imgs_list=[]
    pad_boxes_list=[]
    pad_classes_list=[]

    h_list = [int(s.shape[1]) for s in imgs_list]
    w_list = [int(s.shape[2]) for s in imgs_list]
    max_h = np.array(h_list).max()
    max_w = np.array(w_list).max()
    for i in range(batch_size):
        img=imgs_list[i]
        pad_imgs_list.append(torch.nn.functional.pad(img,(0,int(max_w-img.shape[2]),0,int(max_h-img.shape[1])),value=0.))

    max_num=0
    for i in range(batch_size):
        n=boxes_list[i].shape[0]
        if n>max_num:max_num=n
    for i in range(batch_size):
        pad_boxes_list.append(torch.nn.functional.pad(boxes_list[i],(0,0,0,max_num-boxes_list[i].shape[0]),value=-1))
        pad_classes_list.append(torch.nn.functional.pad(classes_list[i],(0,max_num-classes_list[i].shape[0]),value=-1))


    batch_boxes=torch.stack(pad_boxes_list)
    batch_classes=torch.stack(pad_classes_list)
    batch_imgs=torch.stack(pad_imgs_list)

    return batch_imgs,batch_boxes,batch_classes