训练import os
import argparse
import random


def parse_args():
    """
    args for training.
    """
    parser = argparse.ArgumentParser(description='Parse args for training')
    # for train
    parser.add_argument('--script', type=str, help='training script name')
    parser.add_argument('--config', type=str, default='baseline', help='yaml configure file name')
    parser.add_argument('--save_dir', type=str, help='root directory to save checkpoints, logs, and tensorboard')
    parser.add_argument('--mode', type=str, choices=["single", "multiple", "multi_node"], default="multiple",
                        help="train on single gpu or multiple gpus")
    parser.add_argument('--nproc_per_node', type=int, help="number of GPUs per node", required=True)  # specify when mode is multiple
    parser.add_argument('--use_lmdb', type=int, choices=[0, 1], default=0)  # whether datasets are in lmdb format
    parser.add_argument('--script_prv', type=str, help='training script name')
    parser.add_argument('--config_prv', type=str, default='baseline', help='yaml configure file name')
    parser.add_argument('--use_wandb', type=int, choices=[0, 1], default=0)  # whether to use wandb
    # for knowledge distillation
    parser.add_argument('--distill', type=int, choices=[0, 1], default=0)  # whether to use knowledge distillation
    parser.add_argument('--script_teacher', type=str, help='teacher script name')
    parser.add_argument('--config_teacher', type=str, help='teacher yaml configure file name')

    # for multiple machines
    parser.add_argument('--rank', type=int, help='Rank of the current process.')
    parser.add_argument('--world-size', type=int, help='Number of processes participating in the job.')
    parser.add_argument('--ip', type=str, default='127.0.0.1', help='IP of the current rank 0.')
    parser.add_argument('--port', type=int, default='20000', help='Port of the current rank 0.')

    args = parser.parse_args()

    return args


def main():
    args = parse_args()
    if args.mode == "single":
        train_cmd = "python lib/train/run_training.py --script %s --config %s --save_dir %s --use_lmdb %d " \
                    "--script_prv %s --config_prv %s --distill %d --script_teacher %s --config_teacher %s --use_wandb %d"\
                    % (args.script, args.config, args.save_dir, args.use_lmdb, args.script_prv, args.config_prv,
                       args.distill, args.script_teacher, args.config_teacher, args.use_wandb)
    elif args.mode == "multiple":
        train_cmd = "torchrun --nproc_per_node %d lib/train/run_training.py " \
                    "-script %s --config %s --save_dir %s --use_lmdb %d --script_prv %s --config_prv %s --use_wandb %d " \
                    "-distill %d --script_teacher %s --config_teacher %s" \
                    % (args.nproc_per_node, random.randint(10000, 50000), args.script, args.config, args.save_dir, args.use_lmdb, args.script_prv, args.config_prv, args.use_wandb,
                       args.distill, args.script_teacher, args.config_teacher)
    elif args.mode == "multi_node":
        train_cmd = "torchrun --nproc_per_node %d --nnodes %d --node_rank %d --master_addr %s --master_port %d lib/train/run_training.py " \
                    "-script %s --config %s --save_dir %s --use_lmdb %d --script_prv %s --config_prv %s --use_wandb %d " \
                    "-distill %d --script_teacher %s --config_teacher %s" \
                    % (args.nproc_per_node, args.ip, args.port, args.world_size, args.rank, args.script, args.config, args.save_dir, args.use_lmdb, args.script_prv, args.config_prv, args.use_wandb,
                       args.distill, args.script_teacher, args.config_teacher)
    else:
        raise ValueError("mode should be 'single' or 'multiple'.")
    os.system(train_cmd)


if __name__ == "__main__":
    main()





，报错：torchrun --nproc_per_node=2 tracking/train.py --script ceutrack --config ceutrack_coesot --save_dir ./output --use_wandb 0
[2024-06-17 22:02:03,347] torch.distributed.run: [WARNING] 
[2024-06-17 22:02:03,347] torch.distributed.run: [WARNING] *****************************************
[2024-06-17 22:02:03,347] torch.distributed.run: [WARNING] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
[2024-06-17 22:02:03,347] torch.distributed.run: [WARNING] *****************************************
usage: train.py [-h] [--script SCRIPT] [--config CONFIG] [--save_dir SAVE_DIR] [--mode {single,multiple,multi_node}] --nproc_per_node NPROC_PER_NODE
                [--use_lmdb {0,1}] [--script_prv SCRIPT_PRV] [--config_prv CONFIG_PRV] [--use_wandb {0,1}] [--distill {0,1}] [--script_teacher SCRIPT_TEACHER]
                [--config_teacher CONFIG_TEACHER] [--rank RANK] [--world-size WORLD_SIZE] [--ip IP] [--port PORT]
train.py: error: the following arguments are required: --nproc_per_node
usage: train.py [-h] [--script SCRIPT] [--config CONFIG] [--save_dir SAVE_DIR] [--mode {single,multiple,multi_node}] --nproc_per_node NPROC_PER_NODE
                [--use_lmdb {0,1}] [--script_prv SCRIPT_PRV] [--config_prv CONFIG_PRV] [--use_wandb {0,1}] [--distill {0,1}] [--script_teacher SCRIPT_TEACHER]
                [--config_teacher CONFIG_TEACHER] [--rank RANK] [--world-size WORLD_SIZE] [--ip IP] [--port PORT]
train.py: error: the following arguments are required: --nproc_per_node
[2024-06-17 22:02:08,357] torch.distributed.elastic.multiprocessing.api: [ERROR] failed (exitcode: 2) local_rank: 0 (pid: 305198) of binary: /home/heyuting/anaconda3/envs/mamba/bin/python
Traceback (most recent call last):
  File "/home/heyuting/anaconda3/envs/mamba/bin/torchrun", line 8, in <module>
    sys.exit(main())
  File "/home/heyuting/anaconda3/envs/mamba/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
    return f(*args, **kwargs)
  File "/home/heyuting/anaconda3/envs/mamba/lib/python3.10/site-packages/torch/distributed/run.py", line 806, in main
    run(args)
  File "/home/heyuting/anaconda3/envs/mamba/lib/python3.10/site-packages/torch/distributed/run.py", line 797, in run
    elastic_launch(
  File "/home/heyuting/anaconda3/envs/mamba/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 134, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/home/heyuting/anaconda3/envs/mamba/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 264, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
tracking/train.py FAILED
------------------------------------------------------------
Failures:
[1]:
  time      : 2024-06-17_22:02:08
  host      : server12
  rank      : 1 (local_rank: 1)
  exitcode  : 2 (pid: 305199)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2024-06-17_22:02:08
  host      : server12
  rank      : 0 (local_rank: 0)
  exitcode  : 2 (pid: 305198)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================