Function main

main.py:141–242

(args)

def main(args):
    
    np.random.seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    
    if args.checkpoint_dir is not None:
        pass
    elif args.test_ckpt is not None:
        # if not define the checkpoint-dir, set to the test checkpoint folder as default
        args.checkpoint_dir = os.path.dirname(args.test_ckpt)
        print(f'testing directory: {args.checkpoint_dir}')
    else:
        raise AssertionError(
            'Either checkpoint_dir or test_ckpt should be presented!'
        )
    
    os.makedirs(args.checkpoint_dir, exist_ok=True)
    accelerator = Accelerator()
    set_seed(args.seed)
    
    ### build datasets and dataloaders
    datasets, dataloaders = build_dataset_func(args)
    
    ### build models
    model = build_model_func(args)
    ### set default checkpoint
    checkpoint = None
    
    # testing phase
    if args.test_only:
        
        try:
            checkpoint = torch.load(args.test_ckpt, map_location=torch.device("cpu"))
            model.load_state_dict(checkpoint["model"], strict=False)
        except:
            print('test the model from scratch...')
        
        model, dataloaders['train'], *dataloaders['test'] = accelerator.prepare(
            model, dataloaders['train'], *dataloaders['test']
        )
        
        for test_loader in dataloaders['test']:
            test_loader.dataset.eval_func(
                args,
                -1,
                model,
                accelerator,
                test_loader
            )
        
    # training phase
    else:
        
        assert (
            args.checkpoint_dir is not None
        ), "`--checkpoint_dir` is required to identify the directory to store the checkpoint"
        os.makedirs(args.checkpoint_dir, exist_ok=True)