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:
@ -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:
|
||||
|
Reference in New Issue
Block a user