ultralytics / ultralytics

NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLite
https://docs.ultralytics.com
GNU Affero General Public License v3.0
23.76k stars 4.74k forks source link

How to stop training from freezing and then Crashing my PySide6 Program? #11679

Closed dirtishurt closed 1 week ago

dirtishurt commented 1 week ago

Search before asking

Question

Basically I want to train a model while my Pyside 6 Application is running and whenever I run the model.train() function my windows freeze and close after the model is done training, I would ideally like to avoid this. I've tries using threads, multiprocessing, and threads from PyQt, Is there a way to run the model.train() seperalty from the rest of the program or is this unavoidable?

Additional

No response

github-actions[bot] commented 1 week ago

👋 Hello @dirtishurt, thank you for your interest in Ultralytics YOLOv8 🚀! We recommend a visit to the Docs for new users where you can find many Python and CLI usage examples and where many of the most common questions may already be answered.

If this is a 🐛 Bug Report, please provide a minimum reproducible example to help us debug it.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset image examples and training logs, and verify you are following our Tips for Best Training Results.

Join the vibrant Ultralytics Discord 🎧 community for real-time conversations and collaborations. This platform offers a perfect space to inquire, showcase your work, and connect with fellow Ultralytics users.

Install

Pip install the ultralytics package including all requirements in a Python>=3.8 environment with PyTorch>=1.8.

pip install ultralytics

Environments

YOLOv8 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

Status

Ultralytics CI

If this badge is green, all Ultralytics CI tests are currently passing. CI tests verify correct operation of all YOLOv8 Modes and Tasks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

glenn-jocher commented 1 week ago

@dirtishurt hello! It sounds like you're encountering an issue where your application freezes during training. This typically happens due to the heavy computation involved in training models, which can block the main thread of your PySide6 application.

To tackle this, running the training in a separate process can help. Although you've mentioned attempts with multiprocessing, I recommend specifically using Python's multiprocessing.Process to execute the training. Here is a concise example:

from multiprocessing import Process
from ultralytics import YOLO

def train_model():
    model = YOLO('yolov8n.pt')
    model.train(data='coco128.yaml', epochs=100, imgsz=640)

# Run the training in a separate process
if __name__ == '__main__':
    p = Process(target=train_model)
    p.start()
    p.join()

Ensure this code block runs under the if __name__ == '__main__': condition, especially important on Windows, to avoid recursive subprocess spawning.

This setup should prevent your GUI from freezing. Let me know if this helps or if further assistance is needed! 🚀

dirtishurt commented 1 week ago

I got it working with a QThread() and I'm leaving this comment here if anyone has a similar issue in the future.

from PySide6.QtCore import QThread
from ultralytics import YOLO
class workerThread(QThread):
    def __init__(self, parent):
        super().__init__()
        self.parent = parent
        self.e = 0
        self.model = YOLO(self.parent.model.path)

    def run(self):
        e = 0
        if os.path.exists(os.path.join(self.parent.workingDirectory, 'runs')):
            pass
        else:
            os.mkdir(os.path.join(self.parent.workingDirectory, 'runs'))
        export_path = os.path.join(self.parent.workingDirectory, 'runs')

        if torch.cuda.is_available():
            print('CUDA Compatible Detected, Starting Training')
            self.model.train(data=self.parent.dataset.path, device=0, patience=self.parent.patience,
                             epochs=self.parent.epochs,
                             imgsz=640, project=export_path, verbose=False)

        else:
            print('WARNING CUDA COMPATIBLE GPU NOT DETECTED, TRAINING WILL TAKE LONGER...')
            self.model.train(data=self.parent.dataset.path, device='cpu', patience=self.parent.patience,
                             epochs=self.parent.epochs,
                             imgsz=640, project=export_path, verbose=False)
glenn-jocher commented 1 week ago

@dirtishurt that's great news! 🎉 Thanks for sharing your solution using QThread with the community. It's super helpful for others who might face the same issue and are looking for a proven workaround. Your example is clear and effectively demonstrates how to integrate YOLO training into a PySide6 application without freezing the UI. Great job, and thanks again for contributing! 👍