def __init__(self, model, resume, config, logger_path):
self.config = config
self.device, device_ids = self._prepare_device(config.device)
self.model = model.to(self.device)
if len(device_ids) > 1:
self.model = torch.nn.DataParallel(model, device_ids=device_ids)
self.epochs = config.trainer.epochs
self.save_freq = config.trainer.save_freq
self.verbosity = config.trainer.verbosity
self.checkpoint_dir = config.trainer.checkpoint_dir
mkdir_dir(self.checkpoint_dir)
self.train_logger = Logger(logger_path)
self.monitor = config.trainer.monitor
self.monitor_mode = config.trainer.monitor_mode
assert self.monitor_mode in ['min', 'max', 'off']
self.monitor_best = math.inf if self.monitor_mode == 'min' else -math.inf
self.start_epoch = 1
self.writer = WriterTensorboardX(config.trainer.checkpoint_dir,
self.train_logger,
config.visualization.tensorboardX)
if resume:
self._resume_checkpoint(resume)
type=str,
help='path to checkpoint (default: None)')
parser.add_argument('-d',
'--device',
default="1",
type=str,
help='indices of GPUs to enable (default: all)')
parser.add_argument('-i',
'--input',
default='h36m',
type=str,
help='h36m/input_folder_path]')
parser.add_argument('-o',
'--output',
default='./output',
type=str,
help='Output folder')
parser.add_argument('--interface',
default='openpose',
type=str,
help='2D detection interface')
args = parser.parse_args()
if args.device:
os.environ["CUDA_VISIBLE_DEVICES"] = args.device
if args.resume:
config = torch.load(args.resume)['config']
output_folder = util.mkdir_dir(
'%s/%s' % (args.output, config.trainer.checkpoint_dir.split('/')[-1]))
main(config, args, output_folder)