|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
|
import sys |
|
|
import torch |
|
|
import logging |
|
|
import wandb |
|
|
|
|
|
from utils.arguments import load_opt_command |
|
|
|
|
|
logging.basicConfig(level=logging.INFO) |
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
def init_wandb(args, job_dir, entity='xueyanz', project='xdecoder', job_name='tmp'): |
|
|
wandb_dir = os.path.join(job_dir, 'wandb') |
|
|
os.makedirs(wandb_dir, exist_ok=True) |
|
|
runid = None |
|
|
if os.path.exists(f"{wandb_dir}/runid.txt"): |
|
|
runid = open(f"{wandb_dir}/runid.txt").read() |
|
|
|
|
|
wandb.init(project=project, |
|
|
name=job_name, |
|
|
dir=wandb_dir, |
|
|
entity=entity, |
|
|
resume="allow", |
|
|
id=runid, |
|
|
config={"hierarchical": True},) |
|
|
|
|
|
open(f"{wandb_dir}/runid.txt", 'w').write(wandb.run.id) |
|
|
wandb.config.update({k: args[k] for k in args if k not in wandb.config}) |
|
|
|
|
|
def main(args=None): |
|
|
''' |
|
|
[Main function for the entry point] |
|
|
1. Set environment variables for distributed training. |
|
|
2. Load the config file and set up the trainer. |
|
|
''' |
|
|
|
|
|
opt, cmdline_args = load_opt_command(args) |
|
|
command = cmdline_args.command |
|
|
|
|
|
if cmdline_args.user_dir: |
|
|
absolute_user_dir = os.path.abspath(cmdline_args.user_dir) |
|
|
opt['base_path'] = absolute_user_dir |
|
|
|
|
|
|
|
|
world_size = 1 |
|
|
if 'OMPI_COMM_WORLD_SIZE' in os.environ: |
|
|
world_size = int(os.environ['OMPI_COMM_WORLD_SIZE']) |
|
|
|
|
|
if opt['TRAINER'] == 'xdecoder': |
|
|
from trainer import XDecoder_Trainer as Trainer |
|
|
else: |
|
|
assert False, "The trainer type: {} is not defined!".format(opt['TRAINER']) |
|
|
|
|
|
trainer = Trainer(opt) |
|
|
os.environ['TORCH_DISTRIBUTED_DEBUG']='DETAIL' |
|
|
|
|
|
if command == "train": |
|
|
if opt['rank'] == 0 and opt['WANDB']: |
|
|
wandb.login(key=os.environ['WANDB_KEY']) |
|
|
init_wandb(opt, trainer.save_folder, job_name=trainer.save_folder) |
|
|
trainer.train() |
|
|
elif command == "evaluate": |
|
|
trainer.eval() |
|
|
else: |
|
|
raise ValueError(f"Unknown command: {command}") |
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
|
|
sys.exit(0) |
|
|
|
