Fine Tuning YoloV5 with custom dataset

[Benti98]

Search before asking

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

Question

Hy!
I am working with Yolov5, and i have a question, how to do Fine Tuning with this neural net.
The things that i would is to take the weights of yolov5s neural net and train it to recognize another object in addition to the other 80 classes that already can recognize in order to detect 81 classes.
So train the neural net only with the dataset of the new classes, not with both datasets (coco and custom).

Additional

No response

[glenn-jocher]

@Benti98 hello!

To fine-tune the YOLOv5 neural net for an additional object, you can first download the YOLOv5s pre-trained weights and then train on your custom dataset of the new classes. You can use the --weights flag in the train.py script to specify the location of the pre-trained weights.

Here's an example command to train YOLOv5s on a new dataset without retraining on the COCO dataset:

python train.py --img 640 --batch 16 --epochs 100 --data /path/to/data.yaml --weights /path/to/yolov5s.pt

Make sure to update the data.yaml file to include the new object class and update the number of classes. Also, familiarize yourself with the available command-line arguments and options by running:

python train.py --help

If you run into any specific issues or errors during the training process, feel free to create a new issue and we will be happy to assist you.

[FeliceSchena]

Hello i've similar question, just slightly different.
In my case i've 3 class to train preserving the 0:79 from MsCoco. I have already edited the class id from my labels and created the yaml file. But i'm trying to edit the loss without success.
In my images there may be people and / or objects, these are not labeled. I should mask the loss in case the network predicts an unlabelled coco class. I've already seen the option to weight images, but I don't think it's for me. I just need think that i can multiply the interested cells of loss vector for 0 before the backward. It's possible?
Also I would have to train using a scheduler and job array. If I insert the parameter --resume also to the train 0 is it a problem?
Thanks in advance for your support

[glenn-jocher]

@FeliceSchena hello!

To modify the loss to mask potentially unlabelled COCO classes, you can edit the file ./models/yolo.py and modify the loss function.

You can check the class label associated with each anchor (tile) by accessing the 'tcls' variable. In order to mask the loss for unlabelled classes, you can simply multiply the appropriate cells of the tcls tensor by zero before the backward. Here's an example code snippet to mask the COCO classes:

# Modify this code block inside the `loss` function in yolo.py
...
cls_loss = nn.functional.binary_cross_entropy_with_logits(pred[..., 5:5 + self.nc], tcls[..., :self.nc], reduction='mean', weight=class_weights, )
# Multiply elements of the tensor by zero if the class label is greater than 80 (unlabelled COCO class)
null_mask = tcls[..., :self.nc] > 0.0
cls_loss = cls_loss * null_mask

Regarding the scheduler and job array, you can resume the training with the --resume flag. However, you will need to set the appropriate --batch-size and --epochs parameters for each individual job in the job array.

We hope this helps! If you have any further questions, please don't hesitate to ask.

[FeliceSchena]

Thanks for the quick reply.
I've checked the yolo.py , there isn't the loss function inside. Intestead in loss.py there's a class called compute loss and build targets method did you mean in here?

[glenn-jocher]

@FeliceSchena I apologize for the confusion, you are correct. The loss computation is defined in loss.py and you can modify the compute_loss function to add a mask to the loss.

You can modify the code block for calculating the binary cross-entropy loss to only consider the classes you're interested in by creating a mask. Here's an example snippet of how to modify the compute_loss function in loss.py to add a mask:

def compute_loss(self, p, targets):
    cls_tgts, box_tgts, indices, anchors = self.build_targets(p, targets) # get the indices of the appropriate class

    # apply mask for the specific classes you're interested in (e.g. 80, 81, 82)
    mask = (cls_tgts > 79) & (cls_tgts < 83)  # create a boolean mask to select the appropriate classes
    cls_tgts_masked = torch.where(mask, cls_tgts - 80, torch.zeros_like(cls_tgts))  # shift class labels and set others to 0

    bn, _, na, nt, d = p.shape  # get shape of prediction tensor

    # compute loss using binary cross-entropy with logits and the mask we created
    loss = nn.functional.binary_cross_entropy_with_logits(
        p[..., self.nc:self.nc + 3], cls_tgts_masked[..., None].float(), reduction='none')  # only consider the classes of interest
    loss = (loss * mask.float()[..., None]).mean()

    return loss

Regarding your second question, you can resume training using the --resume flag, and set the batch size and number of epochs accordingly for each iteration in your job array.

Hope this helps! Let me know if you have any further questions.

[FeliceSchena]

@glenn-jocher Thanks for your help. But the p are a list of tensors, so it's impossible to iterate through in this manner:
mask = (cls_tgts > 79) & (cls_tgts < 83)

The same it's for cls_tgts.
I've tryed this code, but the shape of cls_tgts it's not what i'm expecting. Sorry again but it's very difficult to debug such code without documentation about it.

        cls_tgts, box_tgts, indices, anchors = self.build_targets(p, targets) # get the indices of the appropriate class
        loss=[]
        for i in range (self.nl):   
            print(p[i].shape)
            #apply a mask to the class targets to only consider the classes of interest
            mask = torch.logical_and((cls_tgts[i] > 79) ,(cls_tgts[i] < 83))  # create a boolean mask to select the appropriate classes
            cls_tgts_masked = torch.where(mask, cls_tgts[i] - 80, torch.zeros_like(cls_tgts[i]))  # shift class labels and set others to 0
            # compute loss using binary cross-entropy with logits and the mask we created
            loss_tmp=nn.functional.binary_cross_entropy_with_logits(
                p[i][self.nc:self.nc+3], cls_tgts_masked[...,None].float(), reduction='none')  # only consider the classes of interest
            loss.append(loss_tmp * mask.float())

        return mean(loss)

Of this code it's wrong because the shapes doesn't match.

[glenn-jocher]

@FeliceSchena hello! It seems that there was a mistake in my previous reply. Since p is a list of tensors, you will have to iterate through p in order to apply the mask to the appropriate logits. Here's an updated code snippet to modify the loss function in loss.py:

def compute_loss(self, p, targets):
    cls_tgts, box_tgts, indices, anchors = self.build_targets(p, targets)

    mask = (cls_tgts > 79) & (cls_tgts < 83)

    loss = 0
    bn, _, na, nt, d = p[0].shape  # get shape of prediction tensor
    for i in range(self.nl):
        # create boolean mask to select the appropriate classes
        mask_i = mask[:, :, i, :, :]
        # shift class labels and set others to 0
        cls_tgts_i = torch.where(mask_i, cls_tgts - 80, torch.zeros_like(cls_tgts))

        # compute loss using binary cross-entropy with logits and the mask we created
        cls_loss_i = nn.functional.binary_cross_entropy_with_logits(
            p[i][..., self.nc:self.nc + 3], cls_tgts_i[..., None].float(), reduction='none')
        cls_loss_i = (cls_loss_i * mask_i.float()[..., None]).sum() / (mask_i.sum().float() + 1e-16)  # avoid division by 0

        loss += cls_loss_i

    return loss / self.nl

This should calculate the loss across all tensors in p, and apply the mask to the logits for the appropriate classes. I hope this works for you, and thanks for bringing this error to my attention!

[FeliceSchena]

Thank you, but the cls_tgts it's a list too. So i can't use the operator:

cls_tgts > 79

Due to this it's impossible to create a mask of the same dimension of p.

mask = (cls_tgts > 79) & (cls_tgts < 83)

In my case tls_tgls[0] has shape 850

[glenn-jocher]

Hello @FeliceSchena,

I apologize for my previous mistake and any confusion it may have caused. You are correct that cls_tgts is a list, and the approach I suggested will need to be modified to work with this data structure.

One way to create a boolean mask for the appropriate classes is to manually iterate over the classes and use the same criteria for masking the classes as before. Here's an example code snippet to get you started:

def compute_loss(self, p, targets):
    cls_tgts, box_tgts, indices, anchors = self.build_targets(p, targets)
    loss = 0
    for i in range(self.nl):
        # apply a mask to the class targets to only consider the classes of interest
        cls_mask = torch.zeros_like(cls_tgts[i], dtype=torch.bool)  # create a boolean mask tensor
        for cls_idx in range(80, 83):
            cls_mask = (cls_mask | (cls_tgts[i] == cls_idx))  # set mask element to true if the class is in the list of valid classes

        # compute loss using binary cross-entropy with logits and the mask we created
        cls_loss_i = nn.functional.binary_cross_entropy_with_logits(
            p[i][..., self.nc:self.nc+3], cls_tgts[i][..., None].float(), reduction='none')
        cls_loss_i = (cls_loss_i * cls_mask.float()[..., None]).sum() / (cls_mask.sum().float() + 1e-16)  # avoid division by 0

        loss += cls_loss_i

    return loss / self.nl

This code should create a boolean mask for the classes of interest and apply it to the corresponding logits in the p tensor. I hope this helps, and let me know if you have any additional questions or concerns.

[ithmz]

@glenn-jocher it maybe a bit of disrespectful but your replies seems a chatgpt generated content, to me 😅

[FeliceSchena]

The last answers were also incorrect. However with a little patience they guided me to understand what the given shape of the tensor represented.
I arrived at this code that I replaced the one you find on line 154.

 mask=torch.logical_and(tcls[i]>79,tcls[i]<83)
 masked_p=mask.where(mask,1.0)
 masked_p=masked_p.where(~mask,self.cn)
 t = torch.full_like(pcls, self.cn, device=self.device)  # targets
t=t.float()
t[range(n), tcls[i]] = self.cp*masked_p
# compute loss using binary cross-entropy with logits and the mask we created
lcls += (self.BCEcls(pcls, t))

It's definitely not the height of elegance, but they seem to work.
This line:

 masked_p=masked_p.where(~mask,self.cn)

It should be redundant, but instead of 0 I decided to use the negatives on the classes.

Furthermore:

t=t.float()

It allowed me to increase the accuracy of the tensor and therefore be able to multiply the targets with something that could be different than 1 and 0.
Actually i'm still training the network on the whole dataset.
I've tested this line on coco128 and a minibatch of my dataset and seems to work, without losing the previously computed information on coco.
To make this work the dataloader must be balanced on the misc dataset.

[glenn-jocher]

@FeliceSchena hello,

Thank you for providing an update regarding your progress and for sharing your code for solving the issue for your specific use case. It's often the case that different applications of a model require slightly different modifications to the code and it's great to see that you were able to find a solution that worked for you.

If you have any further questions or need any additional assistance, feel free to ask and the community will do its best to help.

Best regards,

YOLOv5 Team.

[github-actions]

👋 Hello there! We wanted to give you a friendly reminder that this issue has not had any recent activity and may be closed soon, but don't worry - you can always reopen it if needed. If you still have any questions or concerns, please feel free to let us know how we can help.

For additional resources and information, please see the links below:

• Docs: https://docs.ultralytics.com
• HUB: https://hub.ultralytics.com
• Community: https://community.ultralytics.com

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLO 🚀 and Vision AI ⭐