Function main

classification/main.py:167–400

(config)

def main(config):
    # prepare data loaders
    dataset_train, dataset_val, dataset_test, data_loader_train, \
    data_loader_val, data_loader_test, mixup_fn = build_loader(config)

    # build runner
    logger.info(f'Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}')
    model = build_model(config)
    model.cuda()
    logger.info(str(model))

    # build optimizer
    optimizer = build_optimizer(config, model)

    if config.AMP_OPT_LEVEL != 'O0':
        config.defrost()
        if has_native_amp:
            config.native_amp = True
            use_amp = 'native'
        elif has_apex:
            config.apex_amp = True
            use_amp = 'apex'
        else:
            use_amp = None
            logger.warning(
                'Neither APEX or native Torch AMP is available, using float32. '
                'Install NVIDA apex or upgrade to PyTorch 1.6')
        config.freeze()

    # setup automatic mixed-precision (AMP) loss scaling and op casting
    amp_autocast = suppress  # do nothing
    loss_scaler = None
    if config.AMP_OPT_LEVEL != 'O0':
        if use_amp == 'apex':
            model, optimizer = amp.initialize(model,
                                              optimizer,
                                              opt_level=config.AMP_OPT_LEVEL)
            loss_scaler = ApexScaler()
            if config.LOCAL_RANK == 0:
                logger.info(
                    'Using NVIDIA APEX AMP. Training in mixed precision.')
        if use_amp == 'native':
            amp_autocast = torch.cuda.amp.autocast
            loss_scaler = NativeScaler()
            if config.LOCAL_RANK == 0:
                logger.info(
                    'Using native Torch AMP. Training in mixed precision.')
        else:
            if config.LOCAL_RANK == 0:
                logger.info('AMP not enabled. Training in float32.')

    # put model on gpus
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False)

    # try:
    #     model.register_comm_hook(state=None, hook=fp16_compress_hook)
    #     logger.info('using fp16_compress_hook!')