AttributeError: 'dict' object has no attribute 'box' when trying to train YOLOv10 on pytoch lightning using v10DetectLoss.

[Vicks0712]

Search before asking

• I have searched the Ultralytics YOLO issues and discussions and found no similar questions.

Question

Hello @THU-MIG team,

I am trying to train the Yolov10 model with a more customized trainer, instead of using the one provided by the default library. I'm trying to reproduce the loss calculation for this model using the v10DetectLoss class, but it always gives this same error:

File ~\AppData\Local\pypoetry\Cache\virtualenvs\hermes-video--0vDOw8A-py3.11\Lib\site-packages\ultralytics\utils\loss.py:724, in v10DetectLoss.__call__(self, preds, batch)
    722 def __call__(self, preds, batch):
    723     one2many = preds["one2many"]
--> 724     loss_one2many = self.one2many(one2many, batch)
    725     one2one = preds["one2one"]
    726     loss_one2one = self.one2one(one2one, batch)

File ~\AppData\Local\pypoetry\Cache\virtualenvs\hermes-video--0vDOw8A-py3.11\Lib\site-packages\ultralytics\utils\loss.py:243, in v8DetectionLoss.__call__(self, preds, batch)
    238     target_bboxes /= stride_tensor
    239     loss[0], loss[2] = self.bbox_loss(
    240         pred_distri, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask
    241     )
--> 243 loss[0] *= self.hyp.box  # box gain
    244 loss[1] *= self.hyp.cls  # cls gain
    245 loss[2] *= self.hyp.dfl  # dfl gain

AttributeError: 'dict' object has no attribute 'box'

And I think it comes from this part of the code in v8DetectionLoss class, in utils/loss.py:

self.hyp = model.args

The code that I'm using is this one, including a custom dataloader and a custom Trainer using pytorch lightining:

dataloader:

class YOLOv10Dataset(Dataset):
    def __init__(self, images_dir, labels_dir, transform=None):
        self.images_dir = images_dir
        self.labels_dir = labels_dir
        self.image_paths = list(Path(images_dir).glob("*.png"))
        self.transform = transform

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

    def __getitem__(self, idx):
        img_path = str(self.image_paths[idx])
        label_path = str(Path(self.labels_dir) / (Path(img_path).stem + ".txt"))

        image = cv2.imread(img_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # Resize and normalize the image
        resized_image = cv2.resize(image, (640, 640))
        normalized_image = resized_image / 255.0
        image_tensor = torch.from_numpy(normalized_image).permute(2, 0, 1).float()

        # Read the labels from the text file
        with open(label_path, 'r') as f:
            labels = np.array([list(map(float, line.strip().split())) for line in f])

        # Extract classes and bounding boxes from the labels
        if len(labels) > 0:
            cls = labels[:, 0]  # First column is the class
            bboxes = labels[:, 1:]  # The other columns are the bounding box coordinates
        else:
            cls = np.array([])  # If there are no labels
            bboxes = np.array([])

        target = {
            'batch_idx': idx,  # Add the batch index
            'cls': torch.tensor(cls, dtype=torch.float32),  # Classes
            'bboxes': torch.tensor(bboxes, dtype=torch.float32)  # Bounding boxes
        }

        return image_tensor, target

train_dataset = YOLOv10Dataset(images_dir="./test/data/images/", labels_dir=Path("./test/data/labels/"))

def collate_fn(batch):
    images = []
    batch_targets = {'batch_idx': [], 'cls': [], 'bboxes': []}

    for i, (image, target) in enumerate(batch):
        images.append(image)

        # Add the data from each image to the corresponding batches
        batch_targets['batch_idx'].append(torch.full((len(target['cls']),), i))  # Batch indices
        batch_targets['cls'].append(target['cls'])
        batch_targets['bboxes'].append(target['bboxes'])

    # Concatenate the data so PyTorch can handle it
    batch_targets['batch_idx'] = torch.cat(batch_targets['batch_idx'])
    batch_targets['cls'] = torch.cat(batch_targets['cls'])
    batch_targets['bboxes'] = torch.cat(batch_targets['bboxes'])

    return torch.stack(images, dim=0), batch_targets


train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, collate_fn=collate_fn)

Loss function (v10DetecLoss):

model = YOLO("yolov10n.pt").to(device)
pytorch_model = model.model.train()
loss_fn = v10DetectLoss(pytorch_model)

Custom Trainer:

class YOLOv10LightningModule(pl.LightningModule):
    def __init__(self, model, lr=1e-3):
        super(YOLOv10LightningModule, self).__init__()
        self.model = model
        self.lr = lr
        self.train_loss = []  #
        self.val_loss = []

    def forward(self, x):
        return self.model(x)

    def compute_loss(self, predictions, targets):
        """
        Calculates the loss across multiple scales from both 'one2many' and 'one2one'.
        """
        loss = loss_fn(predictions, targets)
        return loss

    def training_step(self, batch, batch_idx):
        images, targets = batch

        images = images.to(self.device)
        targets['batch_idx'] = targets['batch_idx'].to(self.device)
        targets['cls'] = targets['cls'].to(self.device)
        targets['bboxes'] = targets['bboxes'].to(self.device)

        predictions = self.forward(images)
        loss = self.compute_loss(predictions, targets)

        self.train_loss.append(loss.item())

        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
        return optimizer

Here are the relevant versions of my dependencies:

PyTorch: 2.4.0
torchvision: 0.19.0
ultralytics: 8.1.34
numpy: 1.26.4

Thank you very much in advance for any kind of help!