text: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import torch
  import torchvision.transforms as transforms
  
  import argparse
  import os
  
  import _init_paths
  from core.config import config
  from core.config import update_config
  from core.config import update_dir
  from utils.utils import create_logger
  import dataset
  import models
  
  def parse_args():
      parser = argparse.ArgumentParser(description='Train keypoints network')
      parser.add_argument(
          '--cfg', help='experiment configure file name', required=True, type=str)
      args, rest = parser.parse_known_args()
      update_config(args.cfg)
  
      parser.add_argument(
          '--frequent',
          help='frequency of logging',
          default=config.PRINT_FREQ,
          type=int)
      parser.add_argument('--gpus', help='gpus', type=str)
      parser.add_argument('--workers', help='num of dataloader workers', type=int)
  
      parser.add_argument(
          '--modelDir', help='model directory', type=str, default='')
      parser.add_argument('--logDir', help='log directory', type=str, default='')
      parser.add_argument(
          '--dataDir', help='data directory', type=str, default='')
      parser.add_argument(
          '--data-format', help='data format', type=str, default='')
      args = parser.parse_args()
      update_dir(args.modelDir, args.logDir, args.dataDir)
      return args
  
  
  def reset_config(config, args):
      if args.gpus:
          config.GPUS = args.gpus
      if args.data_format:
          config.DATASET.DATA_FORMAT = args.data_format
      if args.workers:
          config.WORKERS = args.workers
  
  
  def main():
      args = parse_args()
      reset_config(config, args)
  
      logger, final_output_dir, tb_log_dir = create_logger(
          config, args.cfg, 'train')
  
      gpus = [int(i) for i in config.GPUS.split(',')]
  
      backbone_model = eval('models.' + config.BACKBONE_MODEL + '.get_pose_net')(
          config, is_train=True)
  
      model = eval('models.' + config.MODEL + '.get_multiview_pose_net')(
          backbone_model, config)
  
      if config.TEST.MODEL_FILE:
          logger.info('=> loading model from {}'.format(config.TEST.MODEL_FILE))
          model.load_state_dict(torch.load(config.TEST.MODEL_FILE))
      else:
          model_path = 'model_best.pth.tar' if config.TEST.STATE == 'best' else 'final_state.pth.tar'
          model_state_file = os.path.join(final_output_dir, model_path)
          logger.info('=> loading model from {}'.format(model_state_file))
          model.load_state_dict(torch.load(model_state_file))
  
      # Data loading code
      normalize = transforms.Normalize(
          mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
      train_dataset = eval('dataset.' + config.DATASET.TRAIN_DATASET)(
          config, config.DATASET.TRAIN_SUBSET, True,
          transforms.Compose([
              transforms.ToTensor(),
              normalize,
          ]))
      valid_dataset = eval('dataset.' + config.DATASET.TEST_DATASET)(
          config, config.DATASET.TEST_SUBSET, False,
          transforms.Compose([
              transforms.ToTensor(),
              normalize,
          ]))
  
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=config.TRAIN.BATCH_SIZE * len(gpus),
          shuffle=config.TRAIN.SHUFFLE,
          num_workers=0,
          pin_memory=True)
      valid_loader = torch.utils.data.DataLoader(
          valid_dataset,
          batch_size=config.TEST.BATCH_SIZE * len(gpus),
          shuffle=False,
          num_workers=0,
          pin_memory=True)
  
      for _, (input, target, weight, meta) in enumerate(train_loader):
          print(len(input), len(target), len(weight), len(meta))
          print((input[0].shape), (target[0].shape), (weight[0].shape))
          break
  
      for _, (input, target, weight, meta) in enumerate(valid_loader):
          print(len(input), len(target), len(weight), len(meta))
          print((input[0].shape), (target[0].shape), (weight[0].shape))
          break
  
  if __name__ == '__main__':
      main()
  