RuntimeError:给定组=1,大小为[64,3,3,3]的权重,预期输入[4,4,1024,2048]有3个通道



我的模型是这样的,设计用于处理RGB图像。之前我测试了RGB和灰度图像,结果很好。突然间,它无法使用新的RGB数据集。

class UNet(nn.Module):
def __init__(self):
super(UNet, self).__init__()
self.inc = inconv(3, 64)
self.down1 = down(64, 128)
self.down2 = down(128, 256)
self.down3 = down(256, 512)
self.down4 = down(512, 512)
self.up1 = up(1024, 256)
self.up2 = up(512, 128)
self.up3 = up(256, 64)
self.up4 = up(128, 64)
self.sem_out = outconv(64, 2)
self.ins_out = outconv(64, 70)

数据集加载器:

train_dataset = DataLoaderInstanceSegmentation()
train_dataloader = DataLoader(train_dataset, batch_size=4)

培训:

for epoch in range(5):    
for batched in train_dataloader:
images, sem_labels, ins_labels = batched
images = Variable(images)
sem_labels = Variable(sem_labels)
ins_labels = Variable(ins_labels)
model.zero_grad()
sem_predict, ins_predict = model(images)

我找到了解决方案。删除阿尔法通道代码后工作良好

# data = np.array(Image.open(img_path))
# label_seg = np.array(Image.open(seg_mask_path))
# label_ins = np.array(Image.open(ins_mask_path))
data =  self.to_tensor(Image.open(img_path).convert('RGB'))
label_seg =  self.to_tensor(Image.open(seg_mask_path).convert('L'))
label_ins =  self.to_tensor(Image.open(ins_mask_path).convert('L'))

相关内容

最新更新