Change trainer.py inside engine

[Jessica-hub]

Search before asking

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

Question

Hi, I'm currently trying to use a new optimizer for YOLO training. Can I change the build_optimizer function in trainer,py to add the optimizer? Will this affect the running of the YOLO? If I cannot, can you let me know what should be the best way to add the optimizer?

Additional

No response

[glenn-jocher]

Hello! 👋

Yes, you can modify the build_optimizer function in trainer.py to integrate a new optimizer for your YOLO training. This change should not negatively affect the running of YOLO as long as the new optimizer is implemented correctly and is compatible with the model's requirements.

Here's a basic example of how you might modify the function:

def build_optimizer(self):
    if self.optimizer_name == 'new_optimizer':
        return NewOptimizer(self.model.parameters(), lr=self.lr, ...)
    else:
        # existing optimizers
        return torch.optim.SGD(self.model.parameters(), lr=self.lr, ...)

Just ensure that your new optimizer adheres to the expected interface and that any specific requirements of the optimizer are met. If you encounter any issues, feel free to share them here!

Happy coding! 🚀