Function main

diffportrait360_release/code/train.py:169–486

(args)

def main(args):
    
    # ******************************
    # initialize training
    # ******************************
    # assign rank    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    
    
    args.world_size = int(os.environ['WORLD_SIZE'])
    
    args.local_rank = int(os.environ['LOCAL_RANK'])
    print("local_rank", args.local_rank)
    #import pdb;pdb.set_trace()
    args.rank = int(os.environ['RANK'])
    args.device = torch.device("cuda", args.local_rank)
    args.num_gpu = torch.cuda.device_count()
    args.use_gpu = torch.cuda.is_available() and args.num_gpu > 0
    #seg_model = load_model(args, args.model_path, True, False)
    os.makedirs(args.local_image_dir,exist_ok=True)
    os.makedirs(args.local_log_dir,exist_ok=True)
    if args.rank == 0:
        print(args)

    # initial distribution comminucation
    dist.init_process_group("nccl",  rank=args.rank, world_size=args.world_size)
    torch.backends.cuda.matmul.allow_tf32 = False # it doenst work once not trained in A100 64G gpu
    torch.backends.cudnn.benchmark = True

    # set seed for reproducibility
    set_seed(args.seed)

    # visdom / tensorboard
    if args.rank == 0:
        tb_writer = SummaryWriter(log_dir=args.local_log_dir)
    else:
        tb_writer = None
    
    # ******************************
    # create model
    # ******************************
    #import pdb;pdb.set_trace()
    model = create_model(args.model_config).cpu()

    model.sd_locked = args.sd_locked
    model.only_mid_control = args.only_mid_control
    model.to(args.local_rank)
    #seg_model.to(args.local_rank)
    if args.local_rank == 0:
        print('Total base  parameters {:.02f}M'.format(count_param([model])))
    model_ema = None
    
    # ******************************

    # ******************************
    # load pre-trained models
    # ******************************
    optimizer_state_dict = None