我的模型是这样的,设计用于处理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'))