|
|
from loguru import logger
|
|
|
|
|
|
import torch
|
|
|
import torch.backends.cudnn as cudnn
|
|
|
|
|
|
from yolox.core import Trainer, launch
|
|
|
from yolox.exp import get_exp
|
|
|
|
|
|
import argparse
|
|
|
import random
|
|
|
import warnings
|
|
|
|
|
|
|
|
|
def make_parser():
|
|
|
parser = argparse.ArgumentParser("YOLOX train parser")
|
|
|
parser.add_argument("-expn", "--experiment-name", type=str, default=None)
|
|
|
parser.add_argument("-n", "--name", type=str, default=None, help="model name")
|
|
|
|
|
|
|
|
|
parser.add_argument(
|
|
|
"--dist-backend", default="nccl", type=str, help="distributed backend"
|
|
|
)
|
|
|
parser.add_argument(
|
|
|
"--dist-url",
|
|
|
default=None,
|
|
|
type=str,
|
|
|
help="url used to set up distributed training",
|
|
|
)
|
|
|
parser.add_argument("-b", "--batch-size", type=int, default=64, help="batch size")
|
|
|
parser.add_argument(
|
|
|
"-d", "--devices", default=None, type=int, help="device for training"
|
|
|
)
|
|
|
parser.add_argument(
|
|
|
"--local_rank", default=0, type=int, help="local rank for dist training"
|
|
|
)
|
|
|
parser.add_argument(
|
|
|
"-f",
|
|
|
"--exp_file",
|
|
|
default=None,
|
|
|
type=str,
|
|
|
help="plz input your expriment description file",
|
|
|
)
|
|
|
parser.add_argument(
|
|
|
"--resume", default=False, action="store_true", help="resume training"
|
|
|
)
|
|
|
parser.add_argument("-c", "--ckpt", default=None, type=str, help="checkpoint file")
|
|
|
parser.add_argument(
|
|
|
"-e",
|
|
|
"--start_epoch",
|
|
|
default=None,
|
|
|
type=int,
|
|
|
help="resume training start epoch",
|
|
|
)
|
|
|
parser.add_argument(
|
|
|
"--num_machines", default=1, type=int, help="num of node for training"
|
|
|
)
|
|
|
parser.add_argument(
|
|
|
"--machine_rank", default=0, type=int, help="node rank for multi-node training"
|
|
|
)
|
|
|
parser.add_argument(
|
|
|
"--fp16",
|
|
|
dest="fp16",
|
|
|
default=True,
|
|
|
action="store_true",
|
|
|
help="Adopting mix precision training.",
|
|
|
)
|
|
|
parser.add_argument(
|
|
|
"-o",
|
|
|
"--occupy",
|
|
|
dest="occupy",
|
|
|
default=False,
|
|
|
action="store_true",
|
|
|
help="occupy GPU memory first for training.",
|
|
|
)
|
|
|
parser.add_argument(
|
|
|
"opts",
|
|
|
help="Modify config options using the command-line",
|
|
|
default=None,
|
|
|
nargs=argparse.REMAINDER,
|
|
|
)
|
|
|
return parser
|
|
|
|
|
|
|
|
|
@logger.catch
|
|
|
def main(exp, args):
|
|
|
if exp.seed is not None:
|
|
|
random.seed(exp.seed)
|
|
|
torch.manual_seed(exp.seed)
|
|
|
cudnn.deterministic = True
|
|
|
warnings.warn(
|
|
|
"You have chosen to seed training. This will turn on the CUDNN deterministic setting, "
|
|
|
"which can slow down your training considerably! You may see unexpected behavior "
|
|
|
"when restarting from checkpoints."
|
|
|
)
|
|
|
|
|
|
|
|
|
cudnn.benchmark = True
|
|
|
|
|
|
trainer = Trainer(exp, args)
|
|
|
trainer.train()
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
args = make_parser().parse_args()
|
|
|
exp = get_exp(args.exp_file, args.name)
|
|
|
exp.merge(args.opts)
|
|
|
|
|
|
if not args.experiment_name:
|
|
|
args.experiment_name = exp.exp_name
|
|
|
|
|
|
num_gpu = torch.cuda.device_count() if args.devices is None else args.devices
|
|
|
assert num_gpu <= torch.cuda.device_count()
|
|
|
|
|
|
launch(
|
|
|
main,
|
|
|
num_gpu,
|
|
|
args.num_machines,
|
|
|
args.machine_rank,
|
|
|
backend=args.dist_backend,
|
|
|
dist_url=args.dist_url,
|
|
|
args=(exp, args),
|
|
|
)
|
|
|
|
