我试着写一个分割模型，我对这个主题还很陌生，我已经到了死角。从我试着调试的内容来看，我认为我的掩码批次的形状与预测批次的大小不匹配，因此我得到了以下错误：

---------------------------------------------------------------------------

AssertionError                            Traceback (most recent call last)

<ipython-input-84-abd99309752a> in <module>()
      3 for i in range(EPOCHS):
      4   #train_loss = train_func(trainloader,model,optimizer)
----> 5   valid_loss = eval_func(validloader,model)
      6 
      7   if valid_loss <best_loss:

4 frames

<ipython-input-82-328c759ec537> in eval_func(dataloader, model)
      6       images = images.to(DEVICE)
      7       masks = mask.to(DEVICE)
----> 8       logits, loss = model(images,masks)
      9       total_loss += loss.item()
     10     return total_loss / len(dataloader)

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1108         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110             return forward_call(*input, **kwargs)
   1111         # Do not call functions when jit is used
   1112         full_backward_hooks, non_full_backward_hooks = [], []

<ipython-input-79-567e281ae719> in forward(self, images, masks)
     15     if mask != None:
     16       print(logits.size)
---> 17       return logits, lossF(logits,masks)
     18     return logits

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1108         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110             return forward_call(*input, **kwargs)
   1111         # Do not call functions when jit is used
   1112         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.7/dist-packages/segmentation_models_pytorch/losses/dice.py in forward(self, y_pred, y_true)
     58     def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
     59 
---> 60         assert y_true.size(0) == y_pred.size(0)
     61 
     62         if self.from_logits:

AssertionError:

我不知道如何修复代码中的错误。我在SegmentationDataset中尝试了一些调整，但没有帮助。您可以在下面的代码中找到相关（对我来说）的部分。

import albumentations as A
def get_train_augs():
  return A.Compose([
    #A.Resize(IMG_SIZE,IMG_SIZE, interpolation = cv2.INTER_LINEAR),
    A.RandomCrop(width=IMG_SIZE, height=IMG_SIZE),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=.75)
  ])

def get_val_augs():
  return A.Compose([
    A.RandomCrop(width=IMG_SIZE, height=IMG_SIZE),
  ])

class SegmentationDataset(Dataset):
  def __init__(self,df,augumentations):
    self.df = df
    self.augumentations = augumentations

  def __len__(self):
    return len(self.df)

  def __getitem__(self,idx):
    row = self.df.iloc [idx]

    image_path = row.Images
    mask_path = row.Masks 
    image = cv2.imread(image_path)
    image = cv2.cvtColor(np.float32(image), cv2.COLOR_BGR2RGB)

    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    mask = (mask==33)*1 + (mask==79)*1
    mask = (mask>0)*1
    mask = np.expand_dims(mask, axis=-1)
    
    if self.augumentations:
      data = self.augumentations(image = image, mask = mask)
      image = data['image']
      mask = data['mask']
    
    image = np.transpose(image, (2,0,1)).astype(np.float32)
    mask = np.transpose(mask, (2,0,1)).astype(np.float32)
    image = torch.Tensor(image)/255.0
    mask = torch.Tensor(mask)
    mask = torch.round(torch.Tensor(mask))
    
    return image, mask

class SegmentationModel(nn.Module):
  def __init__(self):
    super(SegmentationModel,self).__init__()

    self.backbone = smp.Unet(
        encoder_name=ENCODER,
        encoder_weights=WEIGHTS,
        in_channels =3,
        classes = 1,
        activation=None
      )
  def forward(self,images, masks= None):
    logits = self.backbone(images)
    lossF = DiceLoss(mode = 'binary')
    if mask != None:
      return logits, lossF(logits,masks)
    return logits

def train_func(dataloader, model,optimizer):
  model.train()
  total_loss = 0.0
  for images, masks in tqdm(dataloader):
    images = images.to(DEVICE)
    masks = mask.to(DEVICE)

    optimizer.zero_grad()
    logits, loss = model(images,masks)
    loss.backward()
    optimizer.step()
    total_loss += loss.item()
    print(mask.size)
  return total_loss / len(dataloader)

  train_loss = train_func(trainloader,model,optimizer)

def eval_func(dataloader, model):
  model.eval()
  total_loss = 0.0
  with torch.no_grad():
    for images, masks in tqdm(dataloader):
      images = images.to(DEVICE)
      masks = mask.to(DEVICE)
      logits, loss = model(images,masks)
      total_loss += loss.item()
    return total_loss / len(dataloader)

Train_Images = os.listdir(os.path.join(os.getcwd(), 'uavid_train/Images'))
for k in range(0,len(Train_Images)): Train_Images[k] = 'uavid_train/Images/' + Train_Images[k]
Train_Labels = os.listdir(os.path.join(os.getcwd(), 'uavid_train/Labels'))
for k in range(0,len(Train_Labels)): Train_Labels[k] = 'uavid_train/Labels/' + Train_Labels[k]
Train_DF = pd.DataFrame([Train_Images, Train_Labels]).T
Train_DF.columns = ['Images', 'Masks']

Val_Images = os.listdir(os.path.join(os.getcwd(), 'uavid_val/Images'))
for k in range(0,len(Val_Images)): Val_Images[k] = 'uavid_val/Images/' + Val_Images[k]
Val_Labels = os.listdir(os.path.join(os.getcwd(), 'uavid_val/Labels'))
for k in range(0,len(Val_Labels)): Val_Labels[k] = 'uavid_val/Labels/' + Val_Labels[k]
Val_DF = pd.DataFrame([ Val_Images, Val_Labels]).T
Val_DF.columns = ['Images', 'Masks']

trainloader = DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True)
validloader = DataLoader(valset, batch_size=len(valset))

model = SegmentationModel()
model.to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr = LR)
best_loss = np.Inf

for i in range(EPOCHS):
  #train_loss = train_func(trainloader,model,optimizer) 
  valid_loss = eval_func(validloader,model)

  if valid_loss <best_loss:
    torch.save(model.state_dict(),"best-model.pt")
    print('SAVED')
    best_loss = valid_loss

  print(f"Epoch :  {i+1} Train Loss : {train_loss} Valid Loss : {valid_loss}")