train.py fix

This commit is contained in:
Alex Bezdieniezhnykh
2024-06-19 03:47:57 +03:00
parent 6c50dd19b7
commit eba8b62db8
+4 -2
View File
@@ -100,13 +100,15 @@ if __name__ == '__main__':
# form_dataset()
# create_yaml()
m = latest_model or 'yolov10x.yaml'
print(f'Initial model: {m}')
model = YOLOv10(latest_model or 'yolov10x.yaml')
folder = f'{prefix}2024-06-18'
yaml = abspath(path.join('datasets', today_folder, 'data.yaml'))
yaml = abspath(path.join('datasets', folder, 'data.yaml'))
results = model.train(data=yaml, epochs=100, batch=10, imgsz=640)
shutil.copy(f'{results.save_dir}/weights/best.pt', latest_model)
shutil.copytree(results.save_dir, f'models/{today_folder}')
shutil.copytree(results.save_dir, f'models/{folder}')
shutil.rmtree('runs')
shutil.rmtree('models/zombobase-latest')