### 实现YOLOv8分类训练的多线程管理与按钮控制 要实现通过按钮控制YOLOv8分类训练的启动和停止，并确保GPU资源正确释放以便重新开始新的训练线程，可以通过以下方法来设计程序结构。 #### 1. 使用 `threading` 进行多线程管理 将 YOLOv8 的训练过程封装在一个独立的线程中，这样可以避免主线程被阻塞。当用户点击“启动”按钮时，创建一个新的线程执行训练；当用户点击“停止”按钮时，终止该线程并清理 GPU 资源。 以下是具体的代码示例： ```python import threading from ultralytics import YOLO import torch from PySide6.QtWidgets import QApplication, QPushButton, QVBoxLayout, QWidget class TrainingManager: def __init__(self): self.thread = None self.stop_event = threading.Event() def start_training(self): if not self.thread or not self.thread.is_alive(): self.stop_event.clear() self.thread = threading.Thread(target=self.run_training, args=(self.stop_event,)) self.thread.start() def stop_training(self): if self.thread and self.thread.is_alive(): self.stop_event.set() # 告知子线程需要退出 self.thread.join() # 等待线程结束 torch.cuda.empty_cache() # 清理 GPU 缓存[^1] @staticmethod def run_training(stop_event): model = YOLO('yolov8n-cls.yaml') try: while not stop_event.is_set(): # 如果未收到停止信号，则继续训练 results = model.train( data='data.yaml', epochs=200, batch=16, imgsz=640, device='cuda' if torch.cuda.is_available() else 'cpu', # 自动选择设备[^2] workers=8, optimizer='SGD' ) if stop_event.is_set(): break except Exception as e: print(f"Training interrupted with error: {e}") finally: torch.cuda.empty_cache() # 确保在任何情况下都清理 GPU 缓存[^1] class MainWindow(QWidget): def __init__(self): super().__init__() self.manager = TrainingManager() layout = QVBoxLayout() self.start_button = QPushButton("Start Training") self.start_button.clicked.connect(self.on_start_clicked) self.stop_button = QPushButton("Stop Training") self.stop_button.clicked.connect(self.on_stop_clicked) layout.addWidget(self.start_button) layout.addWidget(self.stop_button) self.setLayout(layout) def on_start_clicked(self): self.manager.start_training() def on_stop_clicked(self): self.manager.stop_training() if __name__ == "__main__": app = QApplication([]) window = MainWindow() window.show() app.exec() ``` --- #### 2. 关键点解析 - **多线程管理** 使用 Python 的 `threading` 模块将训练逻辑放到单独的线程中运行，从而避免 GUI 主线程被阻塞。通过传递一个共享的 `stop_event` 对象，在主线程中通知子线程何时应该停止运行[^3]。 - **GPU 资源释放** 在训练结束后调用 `torch.cuda.empty_cache()` 来手动清空 GPU 上的缓存，防止内存泄漏影响后续训练任务[^1]。 - **动态设备切换** 利用 `device='cuda' if torch.cuda.is_available() else 'cpu'` 动态判断是否有可用的 GPU 设备，如果没有则回退到 CPU 模式[^2]。 --- ####