Get stuck in deadlock while training

[aliceinland]

🐛 Bug

During the training phase, at a certain epoch, the training get stuck. There are no error reported and, when exiting with brutal force, the only information given is "wait_for_tstate_lock".

To Reproduce

class MaskRCNN(pl.LightningModule):
    
    def __init__(self, hidden_layers=256, n_classes=47, learning_rate = 1e-2):
        super().__init__()
        
        self.hidden_layers = hidden_layers
        self.n_classes = n_classes
        self.learning_rate = learning_rate
        
        self.model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True, progress = False)
        in_features = self.model.roi_heads.box_predictor.cls_score.in_features
        # replace the pre-trained head with a new one
        self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, self.n_classes)
        in_features_mask = self.model.roi_heads.mask_predictor.conv5_mask.in_channels
        self.model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, self.hidden_layers, self.n_classes)
    
    def forward(self, x, y):
        return self.model(x, y)
    
    
    '''def on_after_backward(self):
    # example to inspect gradient information in tensorboard
        if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
            params = self.state_dict()
            for k, v in params.items():
                print(v)
                print(type(v))
                grads = v
                name = k
                self.logger.experiment.add_histogram(tag=name, values=grads, global_step=self.trainer.global_step)'''

        
    def training_step(self, batch, batch_idx):

        #images, targets = batch
        images = batch[0]
        targets = batch[1]
        loss_dict = self.model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        
        self.log('Training general loss:', losses)
        for key, value in loss_dict.items():
            #self.logger.experiment.add_scalar("{}".format(key), value, self.current_epoch)
            self.log(key, value)
            
        return losses
    
    def validation_step(self, batch, batch_idx):
        
        #batch = self.batch_to_device(batch, self.trainer.model.device)
        images, targets = batch
        self.model.train()
        loss_dict = self.model(images, targets)
        self.model.eval()
        prediction = self.model(images)
        prediction = np.array(prediction)
        
        val_losses = sum(loss for loss in loss_dict.values())
        self.log("General_Validation_Loss", val_losses)
        
        #self device per muoverle le cose
        for key, value in loss_dict.items():
            self.log(key, value)
        for k in range(0, len(prediction)):
            #print("Boxes {}".format(prediction[k]['boxes']))
            if len(prediction[k]['boxes']) > 0 and len(prediction[k]['masks']>0):
                nms_res = torchvision.ops.nms(prediction[k]['boxes'], prediction[k]['scores'], 0.3)
                #keep only the elements in the nms_res
                prediction[k]['masks'] = torch.index_select(prediction[k]['masks'], 0, nms_res)
                prediction[k]['boxes'] = torch.index_select(prediction[k]['boxes'], 0, nms_res)
                prediction[k]['labels'] = torch.index_select(prediction[k]['labels'], 0, nms_res)
                prediction[k]['scores'] = torch.index_select(prediction[k]['scores'], 0, nms_res)
                #mAP computation
                #print((targets[k]['boxes']).shape)
                #mAP = calculate_map(targets[0]['boxes'], prediction[0]['boxes'], prediction[0]['scores'], thresh = 0.1)
                mAP = calculate_map(targets[k]['boxes'], prediction[k]['boxes'], prediction[k]['scores'], thresh = 0.1)
                #mAP50 = calculate_map(targets[0]['boxes'], prediction[0]['boxes'], prediction[0]['scores'])
                mAP50 = calculate_map(targets[k]['boxes'], prediction[k]['boxes'], prediction[k]['scores'])
                #mAP75 = calculate_map(targets[0]['boxes'], prediction[0]['boxes'], prediction[0]['scores'], thresh = 0.75)
                mAP75 = calculate_map(targets[k]['boxes'], prediction[k]['boxes'], prediction[k]['scores'], thresh = 0.75)

                self.log("mAP_boxes", mAP)
                self.log("mAP50_boxes", mAP50)
                self.log("mAP75_boxes", mAP75)
        return val_losses
     
     
    def configure_optimizers(self):
        optim = torch.optim.SGD(self.parameters(), lr=self.learning_rate, momentum=0.9, weight_decay=0.0001)
        #optim = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        return optim

  checkpoint_callback = ModelCheckpoint(
          monitor='General_Validation_Loss',
          dirpath='./checkpoints_pruned',
          filename='imaterialist_maskrcnn224x224',
          save_top_k=3,
          mode='min'
      )
    
    tb_logger = pl.loggers.TensorBoardLogger('logs/pruned_imaterialist', name = 'general-224x224', default_hp_metric=False)
    lr_monitor = LearningRateMonitor(logging_interval = 'epoch')
    #n_classes = 46+1
    n_classes = 19
    model = MaskRCNN(hidden_layers = 256, n_classes=n_classes, learning_rate = 0.005)
    data_module = FashionDatasetModule("./dataset/train_splitted", "./dataset/val_splitted", "./dataset/train_pruned.csv", "./dataset/val_pruned.csv", 224, 224)
    trainer = pl.Trainer(max_epochs = 250, gpus = 8, accelerator = 'ddp', plugins=DDPPlugin(find_unused_parameters=False), default_root_dir = './checkpoints_pruned',logger = tb_logger, callbacks=[checkpoint_callback, lr_monitor])
    trainer.fit(model, data_module)
  

Expected behavior

Environment

• CUDA:
  • available: True
  • version: 11.2
• Packages:
  • numpy: 1.19.2
  • pyTorch_debug: False
  • pyTorch_version: 1.8.0a0+52ea372
  • pytorch-lightning: 1.3.8
  • tqdm: 4.53.0
• System:
  • OS: Linux
  • python: 3.8.5

• How you installed PyTorch (conda, pip, source): pip

Additional context

To try to solve this problem I removed the LRScheduler, gradient clipping and early stopping, restarted the docker container but with no results.

[awaelchli]

Hey @aliceinland

Did you forget to submit the FashionDatasetModule? There is also an undefined function calculate_map. Could you also share this code so we can run the example.

[aliceinland]

Hi @awaelchli!
No, the dataset module is in the section before as you can see from the code above!

import pytorch_lightning as pl
import collections
import os
import numpy as np
import pandas as pd
import torchvision
import torch
import torch.utils.data
from PIL import Image, ImageFile
from numba import jit
from torchvision import transforms
import itertools
import math
import natsort
import random
import matplotlib.pyplot as plt
import json
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from FashionDataset import FashionDataset
from utils_maskrcnn_lt import *
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.callbacks import LearningRateMonitor
from pytorch_lightning.plugins import DDPPlugin
from pytorch_lightning.callbacks.early_stopping import EarlyStopping

torch.manual_seed(1708)
np.random.seed(1708)
random.seed(1708)

class FashionDatasetModule(pl.LightningDataModule):
       
        def __init__(self, data_dir_train, data_dir_val, data_train_csv, data_val_csv, height, width, transform = True):
            super().__init__()
            self.data_dir_train = data_dir_train
            self.data_dir_val = data_dir_val
            self.data_train_csv = data_train_csv
            self.data_val_csv = data_val_csv
            self.height = height
            self.width = width
            self.transform = transform

        def collate_fn(self, batch):
            return tuple(zip(*batch))
        
        def setup(self, stage):
            # how to split, etc..
            self.train_dataset = FashionDataset(self.data_dir_train, self.data_train_csv, self.height, self.width, self.transform)
            self.transform = False
            self.val_dataset = FashionDataset(self.data_dir_val, self.data_val_csv, self.height, self.width, self.transform)
            
        def train_dataloader(self):
            return torch.utils.data.DataLoader(self.train_dataset, batch_size = 8, shuffle = True, num_workers = 20, collate_fn=self.collate_fn)
            
        def val_dataloader(self):
            return torch.utils.data.DataLoader(self.val_dataset, batch_size = 8, shuffle = False, num_workers = 20, collate_fn=self.collate_fn)

        #def test_dataloader(self):
            #since the test set contains only images, this will be handeled in another way
            # solamente vivendo scopriremo come 
            #return
        
class MaskRCNN(pl.LightningModule):
    
    def __init__(self, hidden_layers=256, n_classes=47, learning_rate = 1e-2):
        super().__init__()
        
        self.hidden_layers = hidden_layers
        self.n_classes = n_classes
        self.learning_rate = learning_rate
        
        self.model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True, progress = False)
        in_features = self.model.roi_heads.box_predictor.cls_score.in_features
        # replace the pre-trained head with a new one
        self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, self.n_classes)
        in_features_mask = self.model.roi_heads.mask_predictor.conv5_mask.in_channels
        self.model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, self.hidden_layers, self.n_classes)
    
    def forward(self, x, y):
        return self.model(x, y)
    
    
    '''def on_after_backward(self):
    # example to inspect gradient information in tensorboard
        if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
            params = self.state_dict()
            for k, v in params.items():
                print(v)
                print(type(v))
                grads = v
                name = k
                self.logger.experiment.add_histogram(tag=name, values=grads, global_step=self.trainer.global_step)'''

        
    def training_step(self, batch, batch_idx):

        #images, targets = batch
        images = batch[0]
        targets = batch[1]
        loss_dict = self.model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        
        self.log('Training general loss:', losses)
        for key, value in loss_dict.items():
            #self.logger.experiment.add_scalar("{}".format(key), value, self.current_epoch)
            self.log(key, value)
            
        return losses
    
    def validation_step(self, batch, batch_idx):
        
        #batch = self.batch_to_device(batch, self.trainer.model.device)
        images, targets = batch
        self.model.train()
        loss_dict = self.model(images, targets)
        self.model.eval()
        prediction = self.model(images)
        prediction = np.array(prediction)
        
        val_losses = sum(loss for loss in loss_dict.values())
        self.log("General_Validation_Loss", val_losses)
        
        #self device per muoverle le cose
        for key, value in loss_dict.items():
            self.log(key, value)
        for k in range(0, len(prediction)):
            #print("Boxes {}".format(prediction[k]['boxes']))
            if len(prediction[k]['boxes']) > 0 and len(prediction[k]['masks']>0):
                nms_res = torchvision.ops.nms(prediction[k]['boxes'], prediction[k]['scores'], 0.3)
                #keep only the elements in the nms_res
                prediction[k]['masks'] = torch.index_select(prediction[k]['masks'], 0, nms_res)
                prediction[k]['boxes'] = torch.index_select(prediction[k]['boxes'], 0, nms_res)
                prediction[k]['labels'] = torch.index_select(prediction[k]['labels'], 0, nms_res)
                prediction[k]['scores'] = torch.index_select(prediction[k]['scores'], 0, nms_res)
                #mAP computation
                #print((targets[k]['boxes']).shape)
                #mAP = calculate_map(targets[0]['boxes'], prediction[0]['boxes'], prediction[0]['scores'], thresh = 0.1)
                mAP = calculate_map(targets[k]['boxes'], prediction[k]['boxes'], prediction[k]['scores'], thresh = 0.1)
                #mAP50 = calculate_map(targets[0]['boxes'], prediction[0]['boxes'], prediction[0]['scores'])
                mAP50 = calculate_map(targets[k]['boxes'], prediction[k]['boxes'], prediction[k]['scores'])
                #mAP75 = calculate_map(targets[0]['boxes'], prediction[0]['boxes'], prediction[0]['scores'], thresh = 0.75)
                mAP75 = calculate_map(targets[k]['boxes'], prediction[k]['boxes'], prediction[k]['scores'], thresh = 0.75)

                self.log("mAP_boxes", mAP)
                self.log("mAP50_boxes", mAP50)
                self.log("mAP75_boxes", mAP75)

            '''i, j = 0, 0
            t = (targets[k]['labels']).detach().cpu().numpy()
            p = (prediction[k]['labels']).detach().cpu().numpy()

            thP = []
            thP50 = []
            while i in range(len(p)) and j in range(len(t)):
                if t[j] == p[i] and (get_iou(prediction[k]['boxes'][i], targets[k]['boxes'][j]) == True):
                    iou = iou_pytorch(prediction[k]['masks'][i, :], targets[k]['masks'][j, :])
                    thP.append(iou.item())
                    if iou > 0.5: #mAP with iou >0.5
                        thP50.append(iou.item())
                        #print("THP {}".format(iou.item()))
                        #self.log.add_image("Mask Prediction", plot_masks(prediction, targets, iou))

                    p = np.delete(p, i)
                    i = 0
                    j = 0
                elif t[j] == p[i] and (get_iou(prediction[k]['boxes'][i], (targets[k]['boxes'][j])) == False):
                    i+=1
                else:
                    j+=1
                if j == (len(targets[k])-1) or i == (len(prediction[k])-1):
                    p = np.delete(p, 0)
                    i, j = 0, 0

            if len(thP) > 0: 
                self.log("Average Intersection over Union", np.mean(np.asarray(thP)), on_epoch = True)
            if len(thP50) > 0:
                self.log("Average Intersection over Union (w. iou>50)", np.mean(np.asarray(thP50)), on_epoch = True)'''
        
        return val_losses
     
     
    def configure_optimizers(self):
        optim = torch.optim.SGD(self.parameters(), lr=self.learning_rate, momentum=0.9, weight_decay=0.0001)
        #optim = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        return optim
        '''lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, mode = 'min', factor = 0.1, patience = 20)
        
        return {
            'optimizer': optim,
            'lr_scheduler': {
                'scheduler': lr_scheduler,
                'monitor': 'General_Validation_Loss'
            }
        }'''

def main():
    
    checkpoint_callback = ModelCheckpoint(
        monitor='General_Validation_Loss',
        dirpath='./checkpoints_pruned',
        filename='imaterialist_maskrcnn224x224',
        save_top_k=3,
        mode='min'
    )
    
    tb_logger = pl.loggers.TensorBoardLogger('logs/pruned_imaterialist', name = 'general-224x224', default_hp_metric=False)
    lr_monitor = LearningRateMonitor(logging_interval = 'epoch')
    #n_classes = 46+1
    n_classes = 19
    model = MaskRCNN(hidden_layers = 256, n_classes=n_classes, learning_rate = 0.005)
    data_module = FashionDatasetModule("./dataset/train_splitted", "./dataset/val_splitted", "./dataset/train_pruned.csv", "./dataset/val_pruned.csv", 224, 224)
    trainer = pl.Trainer(max_epochs = 250, gpus = 8, accelerator = 'ddp', plugins=DDPPlugin(find_unused_parameters=False), default_root_dir = './checkpoints_pruned',logger = tb_logger, callbacks=[checkpoint_callback, lr_monitor])
    trainer.fit(model, data_module)
    
if __name__ == "__main__":
    main()

The calculate_map function, instead, is in another file included at the beginning (utils_maskrcnn_lt)

from torchvision.ops.boxes import box_iou

def align_coordinates(boxes):
    """Align coordinates (x1,y1) < (x2,y2) to work with torchvision `box_iou` op
    Arguments:
        boxes (Tensor[N,4])
    
    Returns:
        boxes (Tensor[N,4]): aligned box coordinates
    """
    x1y1 = torch.min(boxes[:,:2,],boxes[:, 2:])
    x2y2 = torch.max(boxes[:,:2,],boxes[:, 2:])
    boxes = torch.cat([x1y1,x2y2],dim=1)
    return boxes


def calculate_iou(gt, pr):
    """Calculates the Intersection over Union.

    Arguments:
        gt: (torch.Tensor[N,4]) coordinates of the ground-truth boxes (target)
        pr: (torch.Tensor[M,4]) coordinates of the predicted boxes (predicted)
        form: (str) gt/pred coordinates format
            - pascal_voc: [xmin, ymin, xmax, ymax]
            - coco: [xmin, ymin, w, h]
    Returns:
        iou (Tensor[N, M]): the NxM matrix containing the pairwise
        IoU values for every element in boxes1 and boxes2
    """
    gt = gt.clone()
    pr = pr.clone()

    gt[:,2] = gt[:,0] + gt[:,2]
    gt[:,3] = gt[:,1] + gt[:,3]
    pr[:,2] = pr[:,0] + pr[:,2]
    pr[:,3] = pr[:,1] + pr[:,3]

    gt = align_coordinates(gt)
    pr = align_coordinates(pr)
    
    return box_iou(gt,pr)

def get_mappings(iou_mat):
    #print("Mat shape {}".format(iou_mat.shape))
    mappings = torch.zeros_like(iou_mat)
    gt_count, pr_count = iou_mat.shape
    
    #print("IouMAT: {}".format(iou_mat))
    #print("Target Count: ".format(gt_count))
    #print("Predicted Count: ".format(pr_count))
    #first mapping (max iou for first pred_box)
    if not iou_mat[:,0].eq(0.).all():
        # if not a zero column
        mappings[iou_mat[:,0].argsort()[-1],0] = 1

    for pr_idx in range(1,pr_count):
        # Sum of all the previous mapping columns will let 
        # us know which gt-boxes are already assigned
        not_assigned = torch.logical_not(mappings[:,:pr_idx].sum(1)).long()

        # Considering unassigned gt-boxes for further evaluation 
        targets = not_assigned * iou_mat[:,pr_idx]

        # If no gt-box satisfy the previous conditions
        # for the current pred-box, ignore it (False Positive)
        if targets.eq(0).all():
            continue

        # max-iou from current column after all the filtering
        # will be the pivot element for mapping
        pivot = targets.argsort()[-1]
        mappings[pivot,pr_idx] = 1
    return mappings


def calculate_map(gt_boxes,pr_boxes,scores,thresh=0.5):
    # sorting
    pr_boxes = pr_boxes[scores.argsort().flip(-1)]
    iou_mat = calculate_iou(gt_boxes,pr_boxes) 
    
    # thresholding
    iou_mat = iou_mat.where(iou_mat>thresh,torch.tensor(0.).to(torch.cuda.current_device()))
    
    mappings = get_mappings(iou_mat)
    
    # mAP calculation
    tp = mappings.sum()
    fp = mappings.sum(0).eq(0).sum()
    fn = mappings.sum(1).eq(0).sum()
    mAP = tp / (tp+fp+fn)
    
    return mAP

The code for this function was taken from here.

The Lightning Module Dataset call another class (FashionDataset) defined as below:

import torchvision.transforms as T
import torchvision.transforms.functional as TF
import torch
import pandas as pd
import collections
from tqdm import tqdm
import os
from PIL import Image, ImageFile
import numpy as np
import random
import math 
import albumentations as A
from albumentations.pytorch import ToTensorV2

def rle_decode(mask_rle, shape):
    '''
    mask_rle: run-length as string formated: [start0] [length0] [start1] [length1]... in 1d array
    shape: (height,width) of array to return
    Returns numpy array according to the shape, 1 - mask, 0 - background
    '''

    shape = (int(math.ceil(shape[1])), int(math.ceil(shape[0])))
    s = mask_rle.split()
    # gets starts & lengths 1d arrays
    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0::2], s[1::2])]
    starts -= 1
    # gets ends 1d array
    ends = starts + lengths
    #print("Starts: {}".format(starts))
    #print("Lengths: {}".format(lengths))
    #print("Ends: {}".format(ends))
    # creates blank mask image 1d array
    #img = np.zeros(shape[0] * shape[1], dtype=np.uint8)
    #print("Dot: {}".format(np.dot(shape[0], shape[1])))
    #print("Per: {}".format(shape[0]*shape[1]))
    img = np.zeros(np.dot(shape[0], shape[1]), dtype=np.uint8)
    # sets mark pixles
    for lo, hi in zip(starts, ends):
        img[lo:hi] = 1
    # reshape as a 2d mask image
    return img.reshape(shape).T  # Needed to align to RLE direction

class FashionDataset(torch.utils.data.Dataset):
    def __init__(self, image_dir, df_path, height, width, transform = True):
        self.image_dir = image_dir
        self.df = pd.read_csv(df_path)
        self.transform = transform 
        #l = self.df['ClassId']
        #l = [x+1 for x in l]
        #self.df['ClassId'] = l
        self.height = height
        self.width = width
        self.image_info = collections.defaultdict(dict)
        self.df['CategoryId'] = self.df.ClassId.apply(lambda x: str(x).split("_")[0])
        temp_df = self.df.groupby('ImageId')['EncodedPixels', 'CategoryId'].agg(lambda x: list(x)).reset_index()
        size_df = self.df.groupby('ImageId')['Height', 'Width'].mean().reset_index()
        temp_df = temp_df.merge(size_df, on='ImageId', how='left')
        unique_values = sorted(self.df['ClassId'].unique())
        classes_dict = dict()
        
        for i in range(len(unique_values)):
            classes_dict.update({str(i):str(unique_values[i])})
        #print(classes_dict)
        self.values_indexes = list(classes_dict.values())
        #print(self.values_indexes)
        
        for index, row in tqdm(temp_df.iterrows(), total=len(temp_df)):
            image_id = row['ImageId']
            image_path = os.path.join(self.image_dir, image_id+".jpg")
            self.image_info[index]["image_id"] = image_id
            self.image_info[index]["image_path"] = image_path
            self.image_info[index]["width"] = self.width
            self.image_info[index]["height"] = self.height
            self.image_info[index]["labels"] = row["CategoryId"]
            self.image_info[index]["orig_height"] = row["Height"]
            self.image_info[index]["orig_width"] = row["Width"]
            self.image_info[index]["annotations"] = row["EncodedPixels"]
        #print(self.df.head())
        #print(self.image_info[4643])
    def transformation(self, image, mask, boxes):
        #image = TF.to_pil_image(image)
        #width, height = image.size
        #image = TF.resize(image, (224, 224))
        #mask = TF.resize(mask, (224, 224))
        #boxes[:, [0,2]] = width - boxes[:, [2,0]]
            #boxes = TF.resize(boxes, (224, 224))
        #print(f'INDEX TRANSFORM {idx}')
        #if mask is None:
        #    print("MASK IS NONE")
        #if boxes is None:
        #    print("BOXES IS NONE")
        #if image is None:
        #    print("IMG IS NONE")
        
        if random.random() > 0.5:
            image = TF.hflip(image)
            mask = TF.hflip(mask)

        if random.random() > 0.5:
            p = random.randint(-45, 45)
            image = TF.rotate(image, p)
            mask = TF.rotate(mask, p)
                #print(boxes.shape)
                #boxes  = TF.rotate(boxes, p)

        if random.random() > 0.7:
            image = TF.adjust_brightness(image, 0.3)
            image = TF.adjust_contrast(image, 0.3)
            image = TF.adjust_saturation(image, 0.1)
            image = TF.adjust_hue(image, 0.1)
            
        return image, mask, boxes

    def __getitem__(self, idx):
        # load images ad masks
        img_path = self.image_info[idx]["image_path"]
        img = Image.open(img_path).convert("RGB")
        img = img.resize((self.width, self.height), resample=Image.BILINEAR)

        info = self.image_info[idx]
        mask = np.zeros((len(info['annotations']), self.width, self.height), dtype=np.uint8)
        labels = []
        for m, (annotation, label) in enumerate(zip(info['annotations'], info['labels'])):
            sub_mask = rle_decode(annotation, (info['orig_height'], info['orig_width']))
            sub_mask = Image.fromarray(sub_mask)
            sub_mask = sub_mask.resize((self.width, self.height), resample=Image.BILINEAR)
            mask[m, :, :] = sub_mask
            label = self.values_indexes.index(label)
            labels.append(int(float(label)) + 1)

        num_objs = len(labels)
        boxes = []
        new_labels = []
        new_masks = []

        for i in range(num_objs):
            try:
                pos = np.where(mask[i, :, :])
                xmin = np.min(pos[1])
                xmax = np.max(pos[1])
                ymin = np.min(pos[0])
                ymax = np.max(pos[0])
                if abs(xmax - xmin) >= 20 and abs(ymax - ymin) >= 20:
                    boxes.append([xmin, ymin, xmax, ymax])
                    new_labels.append(labels[i])
                    new_masks.append(mask[i, :, :])
            except ValueError:
                continue

        if len(new_labels) == 0:
            boxes.append([0, 0, 20, 20])
            new_labels.append(0)
            new_masks.append(mask[0, :, :])

        nmx = np.zeros((len(new_masks), self.width, self.height), dtype=np.uint8)
        mask_list = []
        for i, n in enumerate(new_masks):
            nmx[i, :, :] = n
            #print(type(n))
            #print(n.shape)
            mask_list.append(n)
        
        #bbox_labels = []
        #for i in range(boxes):
        #    bbox_labels.append(boxes[i], labels[i])
        #print(type(mask_list[0]))
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.as_tensor(new_labels, dtype=torch.int64)
        masks = torch.as_tensor(nmx, dtype=torch.uint8)

        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        target = {}
        #print(masks)
        #print(boxes)
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd
      
        if self.transform == True:
            try:
                img, masks, boxes = self.transformation(img, masks, boxes)
            except:
                print(f"None value for idx {idx}")
                pass
        img = T.ToTensor()(img)
        
        target['boxes'] = boxes
        target['masks'] = masks

        return img, target

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

[awaelchli]

Hey @aliceinland I saw this was closed. I forgot to follow up here, sorry. Have you found the issue?

[aliceinland]

Hi! The issue was solved with the last update, thank you!