def before_train(self):
logger.info("args: {}".format(self.args))
logger.info("exp value:\n{}".format(self.exp))
# model related init
torch.cuda.set_device(self.local_rank)
model = self.exp.get_model()
logger.info("Model Summary: {}".format(
get_model_info(model, self.exp.test_size)))
model.to(self.device)
# solver related init
self.optimizer = self.exp.get_optimizer(self.args.batch_size)
if self.amp_training:
model, optimizer = amp.initialize(model,
self.optimizer,
opt_level="O1")
# value of epoch will be set in `resume_train`
model = self.resume_train(model)
# data related init
self.no_aug = self.start_epoch >= self.max_epoch - self.exp.no_aug_epochs
self.train_loader = self.exp.get_data_loader(
batch_size=self.args.batch_size,
is_distributed=self.is_distributed,
no_aug=self.no_aug,
)
logger.info("init prefetcher, this might take one minute or less...")
self.prefetcher = DataPrefetcher(self.train_loader)
# max_iter means iters per epoch
self.max_iter = len(self.train_loader)
self.lr_scheduler = self.exp.get_lr_scheduler(
self.exp.basic_lr_per_img * self.args.batch_size, self.max_iter)
if self.args.occupy:
occupy_mem(self.local_rank)
if self.is_distributed:
model = apex.parallel.DistributedDataParallel(model)
# from torch.nn.parallel import DistributedDataParallel as DDP
# model = DDP(model, device_ids=[self.local_rank], broadcast_buffers=False)
if self.use_model_ema:
self.ema_model = ModelEMA(model, 0.9998)
self.ema_model.updates = self.max_iter * self.start_epoch
self.model = model
self.model.train()
self.evaluator = self.exp.get_evaluator(
batch_size=self.args.batch_size,
is_distributed=self.is_distributed)
# Tensorboard logger
if self.rank == 0:
self.tblogger = SummaryWriter(self.file_name)
logger.info("Training start...")
logger.info("\n{}".format(model))
def main(exp, args):
if not args.experiment_name:
args.experiment_name = exp.exp_name
# set environment variables for distributed training
cudnn.benchmark = True
rank = 0
file_name = os.path.join(exp.output_dir, args.experiment_name)
os.makedirs(file_name, exist_ok=True)
if args.save_result:
vis_folder = os.path.join(file_name, 'vis_res')
os.makedirs(vis_folder, exist_ok=True)
setup_logger(
file_name, distributed_rank=rank, filename="demo_log.txt", mode="a"
)
logger.info("Args: {}".format(args))
if args.conf is not None:
exp.test_conf = args.conf
if args.nms is not None:
exp.nmsthre = args.nms
if args.tsize is not None:
exp.test_size = (args.tsize, args.tsize)
model = exp.get_model()
logger.info("Model Summary: {}".format(get_model_info(model, exp.test_size)))
torch.cuda.set_device(rank)
model.cuda(rank)
model.eval()
if not args.trt:
if args.ckpt is None:
ckpt_file = os.path.join(file_name, "best_ckpt.pth.tar")
else:
ckpt_file = args.ckpt
logger.info("loading checkpoint")
loc = "cuda:{}".format(rank)
ckpt = torch.load(ckpt_file, map_location=loc)
# load the model state dict
model.load_state_dict(ckpt["model"])
logger.info("loaded checkpoint done.")
if args.fuse:
logger.info("\tFusing model...")
model = fuse_model(model)
if args.trt:
assert (not args.fuse),\
"TensorRT model is not support model fusing!"
trt_file = os.path.join(file_name, "model_trt.pth")
assert os.path.exists(trt_file), (
"TensorRT model is not found!\n Run python3 tools/trt.py first!"
)
model.head.decode_in_inference = False
decoder = model.head.decode_outputs
logger.info("Using TensorRT to inference")
else:
trt_file = None
decoder = None
predictor = Predictor(model, exp, COCO_CLASSES, trt_file, decoder)
current_time = time.localtime()
if args.demo == 'image':
image_demo(predictor, vis_folder, args.path, current_time, args.save_result)
elif args.demo == 'video' or args.demo == 'webcam':
imageflow_demo(predictor, vis_folder, current_time, args)
def main(exp, args):
if not args.experiment_name:
args.experiment_name = exp.exp_name
output_dir = osp.join(exp.output_dir, args.experiment_name)
os.makedirs(output_dir, exist_ok=True)
if args.save_result:
vis_folder = osp.join(output_dir, "track_vis")
os.makedirs(vis_folder, exist_ok=True)
if args.trt:
args.device = "gpu"
args.device = torch.device("cuda" if args.device == "gpu" else "cpu")
logger.info("Args: {}".format(args))
if args.conf is not None:
exp.test_conf = args.conf
if args.nms is not None:
exp.nmsthre = args.nms
if args.tsize is not None:
exp.test_size = (args.tsize, args.tsize)
model = exp.get_model().to(args.device)
logger.info("Model Summary: {}".format(get_model_info(
model, exp.test_size)))
model.eval()
if not args.trt:
if args.ckpt is None:
ckpt_file = osp.join(output_dir, "best_ckpt.pth.tar")
else:
ckpt_file = args.ckpt
logger.info("loading checkpoint")
ckpt = torch.load(ckpt_file, map_location="cpu")
# load the model state dict
model.load_state_dict(ckpt["model"])
logger.info("loaded checkpoint done.")
if args.fuse:
logger.info("\tFusing model...")
model = fuse_model(model)
if args.fp16:
model = model.half() # to FP16
if args.trt:
assert not args.fuse, "TensorRT model is not support model fusing!"
trt_file = osp.join(output_dir, "model_trt.pth")
assert osp.exists(
trt_file
), "TensorRT model is not found!\n Run python3 tools/trt.py first!"
model.head.decode_in_inference = False
decoder = model.head.decode_outputs
logger.info("Using TensorRT to inference")
else:
trt_file = None
decoder = None
predictor = Predictor(model, exp, trt_file, decoder, args.device,
args.fp16)
current_time = time.localtime()
if args.demo == "image":
image_demo(predictor, vis_folder, current_time, args)
elif args.demo == "video" or args.demo == "webcam":
imageflow_demo(predictor, vis_folder, current_time, args)
def main(exp, args, num_gpu):
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
warnings.warn(
"You have chosen to seed testing. This will turn on the CUDNN deterministic setting, "
)
is_distributed = num_gpu > 1
# set environment variables for distributed training
cudnn.benchmark = True
rank = args.local_rank
# rank = get_local_rank()
file_name = os.path.join(exp.output_dir, args.experiment_name)
if rank == 0:
os.makedirs(file_name, exist_ok=True)
setup_logger(file_name,
distributed_rank=rank,
filename="val_log.txt",
mode="a")
logger.info("Args: {}".format(args))
if args.conf is not None:
exp.test_conf = args.conf
if args.nms is not None:
exp.nmsthre = args.nms
if args.tsize is not None:
exp.test_size = (args.tsize, args.tsize)
model = exp.get_model()
logger.info("Model Summary: {}".format(get_model_info(
model, exp.test_size)))
logger.info("Model Structure:\n{}".format(str(model)))
evaluator = exp.get_evaluator(args.batch_size, is_distributed, args.test)
torch.cuda.set_device(rank)
model.cuda(rank)
model.eval()
if not args.speed and not args.trt:
if args.ckpt is None:
ckpt_file = os.path.join(file_name, "best_ckpt.pth.tar")
else:
ckpt_file = args.ckpt
logger.info("loading checkpoint")
loc = "cuda:{}".format(rank)
ckpt = torch.load(ckpt_file, map_location=loc)
# load the model state dict
model.load_state_dict(ckpt["model"])
logger.info("loaded checkpoint done.")
if is_distributed:
model = DDP(model, device_ids=[rank])
if args.fuse:
logger.info("\tFusing model...")
model = fuse_model(model)
if args.trt:
assert (
not args.fuse and not is_distributed and args.batch_size == 1
), "TensorRT model is not support model fusing and distributed inferencing!"
trt_file = os.path.join(file_name, "model_trt.pth")
assert os.path.exists(
trt_file), "TensorRT model is not found!\n Run tools/trt.py first!"
model.head.decode_in_inference = False
decoder = model.head.decode_outputs
else:
trt_file = None
decoder = None
# start evaluate
*_, summary = evaluator.evaluate(model, is_distributed, args.fp16,
trt_file, decoder, exp.test_size)
logger.info("\n" + summary)
def main():
parser = argparse.ArgumentParser(description="")
parser.add_argument("model", type=str, help="Model to export")
parser.add_argument("-o",
"--output_dir",
type=str,
default="",
help="Output directory")
parser.add_argument('-v',
"--verbose",
action='store_true',
help="Set logging level to INFO")
parser.add_argument('--weights',
type=str,
help="yolo-x weights file (.pth or .pt)")
parser.add_argument(
'--backbone_weights',
type=str,
help="yolo-x weights file, but will be applied only to backbone")
parser.add_argument('--yolox_path',
type=str,
help="Path of yolo-x repository")
parser.add_argument('--num_classes',
type=int,
default=80,
help="Number of classes of the model")
parser.add_argument('--gpu', type=int, help="GPU id to run on GPU")
parser.add_argument('--to_onnx',
action="store_true",
help="Export model to onnx")
parser.add_argument(
'--use_wrapper',
action="store_true",
help=
"In case of onnx export, if this option is present, the model will be wrapped so that its output match dede expectations"
)
parser.add_argument(
'--top_k',
type=int,
default=200,
help="When exporting to onnx, specify maximum returned prediction count"
)
parser.add_argument('--batch_size',
type=int,
default=1,
help="When exporting to onnx, batch size of model")
parser.add_argument(
'--img_width',
type=int,
default=640,
help="Width of the image when exporting with fixed image size")
parser.add_argument(
'--img_height',
type=int,
default=640,
help="Height of the image when exporting with fixed image size")
args = parser.parse_args()
if args.verbose:
logging.basicConfig(level=logging.INFO)
device = "cpu"
if args.gpu:
device = "cuda:%d" % args.gpu
# get yolox model
sys.path.insert(0, args.yolox_path)
import yolox
from yolox.exp import get_exp
from yolox.utils import get_model_info, postprocess, replace_module
from yolox.models.network_blocks import SiLU
exp = get_exp(None, args.model)
exp.num_classes = args.num_classes
logging.info("num_classes == %d" % args.num_classes)
model = exp.get_model()
model.eval()
model.head.decode_in_inference = True
if args.weights:
logging.info("Load weights from %s" % args.weights)
try:
# state_dict
weights = torch.load(args.weights)["model"]
except:
# torchscript
logging.info("Detected torchscript weights")
weights = torch.jit.load(args.weights).state_dict()
weights = {k[6:]: w
for k, w in weights.items()} # skip "model." prefix
model.load_state_dict(weights, strict=True)
elif args.backbone_weights:
logging.info("Load weights from %s" % args.backbone_weights)
weights = torch.load(args.backbone_weights)["model"]
weights = {k: w for k, w in weights.items() if "backbone" in k}
model.load_state_dict(weights, strict=False)
logging.info("Model Summary: {}".format(
get_model_info(model, exp.test_size)))
filename = os.path.join(args.output_dir, args.model)
if args.to_onnx:
model = replace_module(model, nn.SiLU, SiLU)
model = YoloXWrapper_TRT(model,
topk=args.top_k,
raw_output=not args.use_wrapper)
model.to(device)
model.eval()
filename += ".onnx"
example = get_image_input(args.batch_size, args.img_width,
args.img_height)
# XXX: dynamic batch size not supported with wrapper
# XXX: dynamic batch size not yet supported in dede as well
dynamic_axes = None # {"input": {0: "batch"}} if not args.use_wrapper else None
torch.onnx.export(model,
example,
filename,
export_params=True,
verbose=args.verbose,
opset_version=11,
do_constant_folding=True,
input_names=["input"],
output_names=["detection_out", "keep_count"],
dynamic_axes=dynamic_axes)
else:
# wrap model
model = YoloXWrapper(model, args.num_classes, postprocess)
model.to(device)
model.eval()
filename += "_cls" + str(args.num_classes) + ".pt"
script_module = torch.jit.script(model)
logging.info("Save jit model at %s" % filename)
script_module.save(filename)