ultralytics 8.0.45 segment CUDA and DDP callback fixes (#1137)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
2023-02-24 12:43:53 +01:00
committed by GitHub
parent bfc078b32f
commit 3765f4f6d9
8 changed files with 49 additions and 51 deletions

View File

@ -24,8 +24,6 @@ def find_free_network_port() -> int:
def generate_ddp_file(trainer):
import_path = '.'.join(str(trainer.__class__).split('.')[1:-1])
if not trainer.resume:
shutil.rmtree(trainer.save_dir) # remove the save_dir
content = f'''cfg = {vars(trainer.args)} \nif __name__ == "__main__":
from ultralytics.{import_path} import {trainer.__class__.__name__}
@ -43,16 +41,17 @@ def generate_ddp_file(trainer):
def generate_ddp_command(world_size, trainer):
import __main__ # noqa local import to avoid https://github.com/Lightning-AI/lightning/issues/15218
# Get file and args (do not use sys.argv due to security vulnerability)
exclude_args = ['save_dir']
args = [f'{k}={v}' for k, v in vars(trainer.args).items() if k not in exclude_args]
file = generate_ddp_file(trainer) # if argv[0].endswith('yolo') else os.path.abspath(argv[0])
# Build command
import __main__ # local import to avoid https://github.com/Lightning-AI/lightning/issues/15218
file = os.path.abspath(sys.argv[0])
using_cli = not file.endswith('.py')
if not trainer.resume:
shutil.rmtree(trainer.save_dir) # remove the save_dir
if using_cli:
file = generate_ddp_file(trainer)
dist_cmd = 'torch.distributed.run' if TORCH_1_9 else 'torch.distributed.launch'
port = find_free_network_port()
exclude_args = ['save_dir']
args = [f'{k}={v}' for k, v in vars(trainer.args).items() if k not in exclude_args]
cmd = [sys.executable, '-m', dist_cmd, '--nproc_per_node', f'{world_size}', '--master_port', f'{port}', file] + args
return cmd, file