DDP and new dataloader Fix (#95)

Co-authored-by: Laughing-q <1185102784@qq.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Ayush Chaurasia
2022-12-26 20:05:49 +05:30
committed by GitHub
parent 16e3c08883
commit 4fb04be20b
6 changed files with 22 additions and 25 deletions

View File

@ -3,7 +3,8 @@ import shutil
import socket
import sys
import tempfile
import time
from . import USER_CONFIG_DIR
def find_free_network_port() -> int:
@ -23,25 +24,25 @@ def find_free_network_port() -> int:
def generate_ddp_file(trainer):
import_path = '.'.join(str(trainer.__class__).split(".")[1:-1])
# remove the save_dir
shutil.rmtree(trainer.save_dir)
shutil.rmtree(trainer.save_dir) # remove the save_dir
content = f'''overrides = {dict(trainer.args)} \nif __name__ == "__main__":
from ultralytics.{import_path} import {trainer.__class__.__name__}
trainer = {trainer.__class__.__name__}(overrides=overrides)
trainer.train()'''
(USER_CONFIG_DIR / 'DDP').mkdir(exist_ok=True)
with tempfile.NamedTemporaryFile(prefix="_temp_",
suffix=f"{id(trainer)}.py",
mode="w+",
encoding='utf-8',
dir=os.path.curdir,
dir=USER_CONFIG_DIR / 'DDP',
delete=False) as file:
file.write(content)
return file.name
def generate_ddp_command(world_size, trainer):
import __main__ # local import to avoid https://github.com/Lightning-AI/lightning/issues/15218
import __main__ # noqa local import to avoid https://github.com/Lightning-AI/lightning/issues/15218
file_name = os.path.abspath(sys.argv[0])
using_cli = not file_name.endswith(".py")
if using_cli:
@ -52,9 +53,7 @@ def generate_ddp_command(world_size, trainer):
def ddp_cleanup(command, trainer):
# delete temp file if created
# TODO: this is a temp solution in case the file is deleted before DDP launching
time.sleep(5)
# delete temp file if created
tempfile_suffix = f"{id(trainer)}.py"
if tempfile_suffix in "".join(command):
for chunk in command: