def main(): logger = get_logger() global_config = config['Global'] # 初始化设备 use_gpu = global_config['use_gpu'] if global_config['local_rank'] == -1 or not use_gpu: device = torch.device( "cuda" if torch.cuda.is_available() and use_gpu else "cpu") global_config.update( {'n_gpu': torch.cuda.device_count() if use_gpu else 1}) else: torch.cuda.set_device(global_config['local_rank']) device = torch.device('cuda', global_config['local_rank']) dist.init_process_group(backend='nccl') global_config.update({'n_gpu': 1}) global_config.update({'device': device}) logger.warning( f"\n\tProcess Rank:{global_config['local_rank']} \n" f"\tDevice: {device}\n" f"\tGpus: {global_config['n_gpu']}\n" f"\tDistributed: {bool(global_config['local_rank'] != -1)}\n" f"\t16-bits training: {global_config['fp16']}") rank_id = global_config['local_rank'] set_seed(global_config['seed'], use_gpu) # 阻塞子进程,下面的操作仅主进程进行 if not is_main_process(rank_id): dist.barrier() post_process = build_post_process(config['PostProcess'], global_config) # 构建模型 arch_config = config.pop('Architecture') if hasattr(post_process, 'character'): char_num = len(getattr(post_process, 'character')) arch_config["Head"]['out_channels'] = char_num logger.info(f"\nModel Info:" f"\n{json.dumps(arch_config, indent=4)}") model = build_model(arch_config) state_dict = torch.load(global_config['pretrained_model']) model.load_state_dict(state_dict) # 加载训练数据 if global_config['local_rank'] == 0: dist.barrier() logger.info(f"\nLoad train Data:" f"\n{json.dumps(config['Train'], indent=4)}") train_dataloader = build_dataloader(config, logger, 'Train') logger.info(f"\nLoad Eval Data:" f"\n{json.dumps(config['Eval'], indent=4)}") eval_dataloader = build_dataloader(config, logger, 'Eval') if global_config['local_rank'] == 0: dist.barrier() model.to(device)
def main(): global_config = config['Global'] device = torch.device('cpu') if global_config['use_gpu'] and torch.cuda.is_available(): device = torch.device('cuda') logger.info('使用设备:{}'.format(device)) logger.info('模型信息:{}'.format(config['Architecture'])) model = build_model(config['Architecture']) model.to(device) logger.info('加载预训练模型:{}'.format(global_config['pretrained_model'])) state_dict = torch.load(global_config['pretrained_model']) model.load_state_dict(state_dict) post_process_class = build_post_process(config['PostProcess']) ops = [] for op in config['Eval']['dataset']['transforms']: op_name = list(op)[0] if 'Label' in op_name: continue elif op_name == "KeepKeys": op[op_name]['keep_keys'] = ['image', 'shape'] ops.append(op) transforms = create_transformers(ops) save_res_path = global_config['save_res_path'] save_dir = os.path.dirname(save_res_path) if not os.path.exists(save_dir): os.makedirs(save_dir) model.eval() with open(save_res_path, 'wb') as fout: for file in get_img_list(global_config['infer_img']): logger.info(f"测试图像:{file}") data = {'image': file} batch = transforms(data) images = np.expand_dims(batch[0], axis=0) shape_list = np.expand_dims(batch[1], axis=0) images = torch.from_numpy(images).to(device) preds = model(images) post_result = post_process_class(preds, shape_list) boxes = post_result[0]['points'] dt_boxes_json = [] for box in boxes: tmp_json = {"transcription": ""} tmp_json['points'] = box.tolist() dt_boxes_json.append(tmp_json) otstr = file + "\t" + json.dumps(dt_boxes_json) + '\n' fout.write(otstr.encode()) src_img = cv.imread(file) draw_det_res(boxes, save_dir, src_img, file) logger.info("结果已保存!")
def main(): global_config = config['Global'] device = torch.device('cpu') if global_config['use_gpu'] is True and torch.cuda.is_available(): device = torch.device('cuda') logger.info('使用设备:{}'.format(device)) post_process_class = build_post_process(config['PostProcess'], global_config) if hasattr(post_process_class, 'character'): config['Architecture']["Head"]['out_channels'] = len( getattr(post_process_class, 'character')) logger.info('构建模型,字典包含{}个字'.format( config['Architecture']["Head"]['out_channels'])) logger.info('模型结构:{}'.format(config['Architecture'])) model = build_model(config['Architecture']) model.to(device) logger.info('加载预训练模型 {}...'.format(global_config['pretrained_model'])) state_dict = torch.load(global_config['pretrained_model']) model.load_state_dict(state_dict) ops = [] for op in config['Eval']['dataset']['transforms']: op_name = list(op)[0] if 'Label' in op_name: continue elif op_name in ['RecResizeImg']: op[op_name]['infer_mode'] = True elif op_name == 'KeepKeys': op[op_name]['keep_keys'] = ['image'] ops.append(op) global_config['infer_mode'] = True transforms = create_transformers(ops, global_config) model.eval() for file in get_img_list(config['Global']['infer_img']): logger.info('输入图像:{}'.format(file)) data = {'image': file} batch = transforms(data) images = torch.from_numpy(batch[0]).to(device) images = images.unsqueeze(0) preds = model(images) post_result = post_process_class(preds) logger.info("result: {}".format(post_result))
def main(): global_config = config['Global'] use_gpu = global_config['use_gpu'] n_gpus = 1 device = torch.device('cpu') if use_gpu: if torch.cuda.is_available(): n_gpus = torch.cuda.device_count() device = torch.device('cuda') else: logger.warning("未发现可用于计算的GPU设备") # 创建数据集 config['Eval']['loader'].update({ 'batch_size': config['Eval']['loader']['batch_size_per_card'] * n_gpus }) dataloader = build_dataloader(config, device, logger, 'Eval') batch_size = config['Eval']['loader']['batch_size'] logger.info(f'测试数据共 {len(dataloader)}个batch, 每个batch包含{batch_size}个样本') post_process_class = build_post_process(config['PostProcess'], global_config) if hasattr(post_process_class, 'character'): config['Architecture']["Head"]['out_channels'] = len( getattr(post_process_class, 'character')) model = build_model(config['Architecture']) # 加载预训练模型 state_dict = torch.load(global_config['pretrained_model'], map_location=torch.device('cpu')) model.load_state_dict(state_dict) model.to(device) eval_class = build_metric(config['Metric']) metric = train_utils.eval(model, dataloader, post_process_class, eval_class, device) logger.info('metric eval ***************') for k, v in metric.items(): logger.info('{}:{}'.format(k, v))
# -*- coding: utf-8 -*- # __author__:Song Zhenqi # 2021-01-20 from argparse import ArgumentParser import paddle from utils import get_logger, get_config from models.architectures import build_model import torch def get_args(): parser = ArgumentParser() parser.add_argument('--paddle', '-p', help='Paddle Model path') parser.add_argument('--cfg', '-c', help='Config File') return parser.parse_args() if __name__ == '__main__': args = get_args() print(args.cfg) config = get_config(args.cfg) config['Architecture']["Head"]['out_channels'] = 6625 net = build_model(config['Architecture']) # static_dict = torch.load('./test.pth') paddle_dict = paddle.load(args.paddle) # net.load_state_dict(static_dict) net.load_paddle_state_dict(paddle_dict) torch.save(net.state_dict(), 'mobilev3_crnn_ctc.pth')