AttributeError: 'dict' object has no attribute 'box' when trying to train YOLOv8 on custom trainer using v8DetectionLoss.

Vicks0712 commented 1 month ago

Search before asking

[X] I have searched the Ultralytics YOLO issues and found no similar bug report.

Ultralytics YOLO Component

No response

Bug

Hello https://github.com/ultralytics team,

I am trying to train the Yolov8 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 v8DetectionLoss 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: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

This is how I'm loading the yolov8 detection model and the loss function:

model = YOLO("yolov8n.pt").to(device)
pytorch_model = model.model.train()
loss_fn = v8DetectionLoss(pytorch_model)

Environment

Ultralytics YOLOv8.1.34 🚀 Python-3.11.8 torch-2.4.0+cu124 CUDA:0 (NVIDIA GeForce GTX 1650, 4096MiB) Setup complete ✅ (12 CPUs, 15.9 GB RAM, 719.3/952.6 GB disk)

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

Minimal Reproducible Example

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = YOLO("yolov8n.pt").to(device)
pytorch_model = model.model.train()

Dataset.py

class YOLOv8Dataset(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 = YOLOv8Dataset(images_dir="./tests/data/images/", labels_dir=Path("./tests/data/labels/"))

Dataloader.py:

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)

Trainer.py:

loss_fn = v8DetectionLoss(pytorch_model)

class YOLOv8LightningModule(pl.LightningModule):
    def __init__(self, model, lr=1e-3):
        super(YOLOv8LightningModule, 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 validation_step(self, batch, batch_idx):
        images, targets = batch
        predictions = self.forward(images)
        loss = self.compute_loss(predictions, targets)

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

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

yolov8_lightning_module = YOLOv8LightningModule(model=pytorch_model, lr=1e-3)
trainer = pl.Trainer(max_epochs=50)
trainer.fit(yolov8_lightning_module, train_loader)

Additional

"Many thanks in advance to the @ultralytics team!!!"

Are you willing to submit a PR?

[ ] Yes I'd like to help by submitting a PR!

UltralyticsAssistant commented 1 month ago

👋 Hello @Vicks0712, thank you for reaching out and bringing this to our attention 🚀!

We suggest checking the Ultralytics Docs for a wealth of examples and guidance, especially around custom training and usage with different components.

If this is a 🐛 Bug Report, please ensure to provide a detailed minimum reproducible example (MRE), which you've done a great job with! This will help us further diagnose the issue.

For custom training ❓ questions, ensure that your dataset and the model's configuration align with our Tips for Best Training Results.

Join the vibrant Ultralytics community for real-time discussions on Discord 🎧, explore detailed topics on Discourse, or engage on our Subreddit.

Upgrade

Make sure your ultralytics package is up-to-date, alongside all requirements, within a Python>=3.8 environment and using PyTorch>=1.8.

pip install -U ultralytics

Environments

Run your model in any of the following environments with all dependencies preinstalled:

Notebooks:
Google Cloud: Follow GCP Quickstart Guide
Amazon: Follow AWS Quickstart Guide
Docker: Follow Docker Quickstart Guide

Status

A green status indicates all Ultralytics CI tests are passing, ensuring correct operations across MAC, Windows, and Ubuntu.

This is an automated response, but rest assured, an Ultralytics engineer will assist you shortly. Thanks for your patience! 😊

Y-T-G commented 1 month ago

Marking as duplicate of https://github.com/ultralytics/ultralytics/issues/16189

You can use the existing issue. Thanks.

ultralytics / ultralytics