Function eval

classification/main_deepspeed.py:457–488

(config)

def eval(config):
    _, _, _, _, data_loader_val, _, _ = build_loader(config)
    model = build_model(config)
    model.cuda()
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False)

    model_wo_ddp = model.module
    if config.MODEL.RESUME:
        try:
            checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu')
            msg = model_wo_ddp.load_state_dict(checkpoint['model'], strict=False)
            logger.info(msg)
        except:
            try:
                from deepspeed.utils.zero_to_fp32 import \
                    get_fp32_state_dict_from_zero_checkpoint
                ckpt_dir = os.path.dirname(config.MODEL.RESUME)
                tag = os.path.basename(config.MODEL.RESUME)
                state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir=ckpt_dir, tag=tag)
                model_wo_ddp.load_state_dict(state_dict)
            except:
                checkpoint = torch.load(os.path.join(config.MODEL.RESUME, 'mp_rank_00_model_states.pt'),
                                        map_location='cpu')
                model_wo_ddp.load_state_dict(checkpoint['module'])
    elif config.MODEL.PRETRAINED:
        load_pretrained(config, model_wo_ddp, logger)

    if config.THROUGHPUT_MODE:
        throughput(data_loader_val, model, logger)

    eval_epoch(config, data_loader_val, model)


if __name__ == '__main__':