def __init__(self, model_path, post_p_thre=0.7, gpu_id=None): ''' 初始化pytorch模型 :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件) model_path='/home/share/gaoluoluo/dbnet/output/DBNet_resnet18_FPN_DBHead/checkpoint/model_latest.pth' :param gpu_id: 在哪一块gpu上运行 ''' self.gpu_id = gpu_id if self.gpu_id is not None and isinstance( self.gpu_id, int) and torch.cuda.is_available(): self.device = torch.device("cuda:%s" % self.gpu_id) else: self.device = torch.device("cpu") # print('device:', self.device) checkpoint = torch.load(model_path, map_location=self.device) # print("checkpoint:",checkpoint) config = checkpoint['config'] # print(checkpoint['config']) config['arch']['backbone']['pretrained'] = False self.model = build_model(config['arch']) self.post_process = get_post_processing(config['post_processing']) self.post_process.box_thresh = post_p_thre self.img_mode = config['dataset']['train']['dataset']['args'][ 'img_mode'] self.model.load_state_dict(checkpoint['state_dict']) self.model.to(self.device) self.model.eval() self.transform = [] for t in config['dataset']['train']['dataset']['args']['transforms']: if t['type'] in ['ToTensor', 'Normalize']: self.transform.append(t) self.transform = get_transforms(self.transform)
def __init__(self, model_path, gpu_id=None): ''' 初始化pytorch模型 :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件) :param gpu_id: 在哪一块gpu上运行 ''' self.gpu_id = gpu_id if self.gpu_id is not None and isinstance(self.gpu_id, int) and torch.cuda.is_available(): self.device = torch.device("cuda:%s" % self.gpu_id) else: self.device = torch.device("cpu") print('device:', self.device) checkpoint = torch.load(model_path, map_location=self.device) config = checkpoint['config'] config['arch']['args']['pretrained'] = False self.model = get_model(config['arch']) self.post_process = get_post_processing(config['post_processing']) self.img_mode = config['dataset']['train']['dataset']['args']['img_mode'] self.model.load_state_dict(checkpoint['state_dict']) self.model.to(self.device) self.model.eval() self.transform = [] for t in config['dataset']['train']['dataset']['args']['transforms']: if t['type'] in ['ToTensor', 'Normalize']: self.transform.append(t) self.transform = get_transforms(self.transform)
def __init__(self, model_path, gpu_id=0): from models import build_model from data_loader import get_dataloader from post_processing import get_post_processing from utils import get_metric self.gpu_id = gpu_id if self.gpu_id is not None and isinstance( self.gpu_id, int) and torch.cuda.is_available(): self.device = torch.device("cuda:%s" % self.gpu_id) torch.backends.cudnn.benchmark = True else: self.device = torch.device("cpu") print('load model:', model_path) checkpoint = torch.load(model_path, map_location=torch.device('cpu')) config = checkpoint['config'] config['arch']['backbone']['pretrained'] = False self.validate_loader = get_dataloader(config['dataset']['validate'], config['distributed']) self.model = build_model(config['arch'].pop('type'), **config['arch']) self.model.load_state_dict(checkpoint['state_dict']) self.model.to(self.device) self.post_process = get_post_processing(config['post_processing']) self.metric_cls = get_metric(config['metric'])
def __init__(self, model_path, gpu_id=0): from models import build_model from data_loader import get_dataloader from post_processing import get_post_processing from utils import get_metric self.gpu_id = gpu_id if self.gpu_id is not None and isinstance(self.gpu_id, int) and torch.cuda.is_available(): self.device = torch.device("cuda:%s" % self.gpu_id) torch.backends.cudnn.benchmark = True else: self.device = torch.device("cpu") # print(self.gpu_id) 0 checkpoint = torch.load(model_path, map_location=torch.device('cpu')) config = checkpoint['config'] config['arch']['backbone']['pretrained'] = False config['dataset']['train']['dataset']['args']['data_path'][0] = '/home/share/gaoluoluo/dbnet/datasets/train_zhen.txt' config['dataset']['validate']['dataset']['args']['data_path'][0] = '/home/share/gaoluoluo/dbnet/datasets/test_zhen.txt' print("config:",config) self.validate_loader = get_dataloader(config['dataset']['validate'], config['distributed']) self.model = build_model(config['arch']) self.model.load_state_dict(checkpoint['state_dict']) self.model.to(self.device) self.post_process = get_post_processing(config['post_processing']) self.metric_cls = get_metric(config['metric'])
def __init__(self, model_path, post_p_thre=0.7, gpu_id=None, save_wts=False): ''' 初始化pytorch模型 :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件) :param gpu_id: 在哪一块gpu上运行 ''' self.gpu_id = gpu_id if self.gpu_id is not None and isinstance( self.gpu_id, int) and torch.cuda.is_available(): self.device = torch.device("cuda:%s" % self.gpu_id) else: self.device = torch.device("cpu") print('device:', self.device) checkpoint = torch.load(model_path, map_location=self.device) config = checkpoint['config'] config['arch']['backbone']['pretrained'] = False self.model = build_model(config['arch']) self.post_process = get_post_processing(config['post_processing']) self.post_process.box_thresh = post_p_thre self.img_mode = config['dataset']['train']['dataset']['args'][ 'img_mode'] self.model.load_state_dict(checkpoint['state_dict']) self.model.to(self.device) self.model.eval() # 保存wts # save wts if save_wts: f = open('DBNet.wts', 'w') f.write('{}\n'.format(len(self.model.state_dict().keys()))) for k, v in self.model.state_dict().items(): vr = v.reshape(-1).cpu().numpy() f.write('{} {} '.format(k, len(vr))) for vv in vr: f.write(' ') f.write(struct.pack('>f', float(vv)).hex()) f.write('\n') self.transform = [] for t in config['dataset']['train']['dataset']['args']['transforms']: if t['type'] in ['ToTensor', 'Normalize']: self.transform.append(t) self.transform = get_transforms(self.transform) self.init_onnx_session() self.init_ie()
def main_entrance(): os.environ["CUDA_VISIBLE_DEVICES"] = '3' args = init_args() config = anyconfig.load(open(args.config_file, 'rb')) # print('===config:', config) if 'base' in config: config = parse_config(config) print('===config:', config) if torch.cuda.device_count() > 1: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend="nccl", init_method="env://", world_size=torch.cuda.device_count(), rank=args.local_rank) config['distributed'] = True else: config['distributed'] = False config['local_rank'] = args.local_rank logging.info(config['dataset']['train']) model = build_model(config['arch']['type'], **config['arch']) criterion = build_loss(config['loss'].pop('type'), **config['loss']).cuda() post_process = get_post_processing(config['post_processing']) train_loader = get_trainloader(dataset.ICDAR2015Dataset, config) eval_loader = get_evalloader(dataset.ICDAR2015Dataset, config) model = model.cuda() if config['distributed']: model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, broadcast_buffers=False, find_unused_parameters=True) checkpoint_path = config['train']['resume_checkpoint'] output_path = config['train']['output_path'] optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=0.00005) # load_weights(model, optimizer, config['distributed'], checkpoint_path='/red_detection/DBNet/code_pretrain_model/model_latest_express_code_7_13.pth') # load_weights(model, optimizer, config['distributed'], checkpoint_path=checkpoint_path) epochs = config['train']['epochs'] warmup_iters = config['lr_scheduler']['args']['warmup_epoch']*len(train_loader) # scheduler = WarmupPolyLR(optimizer, max_iters=epochs * len(train_loader), # warmup_iters=warmup_iters, **config['lr_scheduler']['args']) train(model, optimizer, epochs, criterion, train_loader, config, post_process, eval_loader,output_path) from matplotlib import pyplot as plt plt.plot(lr_list) plt.savefig('./show_lr_word_industry.png')
def main(config): import torch from models import build_model, build_loss from data_loader import get_dataloader from trainer import Trainer from post_processing import get_post_processing from utils import get_metric if torch.cuda.device_count() > 1: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group( backend="nccl", init_method="env://", world_size=torch.cuda.device_count(), rank=args.local_rank) config['distributed'] = True else: config['distributed'] = False config['local_rank'] = args.local_rank train_loader = get_dataloader(config['dataset']['train'], config['distributed']) assert train_loader is not None if 'validate' in config['dataset']: validate_loader = get_dataloader(config['dataset']['validate'], False) else: validate_loader = None criterion = build_loss(config['loss'].pop('type'), **config['loss']).cuda() config['arch']['backbone']['in_channels'] = 3 if config['dataset'][ 'train']['dataset']['args']['img_mode'] != 'GRAY' else 1 config['arch']['backbone']['pretrained'] = False model = build_model(config['arch']['type'], **config['arch']) post_p = get_post_processing(config['post_processing']) metric = get_metric(config['metric']) trainer = Trainer(config=config, model=model, criterion=criterion, train_loader=train_loader, post_process=post_p, metric_cls=metric, validate_loader=validate_loader) trainer.train()
def __init__(self, model_path, gpu_id=0): from models import get_model from data_loader import get_dataloader from post_processing import get_post_processing from utils import get_metric self.device = torch.device("cuda:%s" % gpu_id) if gpu_id is not None: torch.backends.cudnn.benchmark = True checkpoint = torch.load(model_path, map_location=torch.device('cpu')) config = checkpoint['config'] config['arch']['args']['pretrained'] = False self.validate_loader = get_dataloader(config['dataset']['validate'], config['distributed']) self.model = get_model(config['arch']) self.model.load_state_dict(checkpoint['state_dict']) self.model.to(self.device) self.post_process = get_post_processing(config['post_processing']) self.metric_cls = get_metric(config['metric'])
def __init__(self, model_path, post_p_thre=0.7): checkpoint = torch.load(model_path) model_config = { 'backbone': { 'type': 'resnet18', 'pretrained': False, "in_channels": 3 }, 'neck': { 'type': 'FPN', 'inner_channels': 256 }, # 分割头,FPN or FPEM_FFM 'head': { 'type': 'DBHead', 'out_channels': 2, 'k': 50 }, } self.model = build_model('Model', **model_config).cuda() self.post_process = get_post_processing({ 'type': 'SegDetectorRepresenter', 'args': { 'thresh': 0.5, 'box_thresh': 0.7, 'max_candidates': 1000, 'unclip_ratio': 1.7 } }) self.post_process.box_thresh = post_p_thre self.img_mode = 'RGB' self.model.load_state_dict(checkpoint) self.model.eval() self.transform = [] self.transform = transforms.Compose([ transforms.ToTensor(), # 转为Tensor 归一化至0~1 transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), # 归一化 ])
def __init__(self, model_path, post_p_thre=0.7): ''' 初始化pytorch模型 :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件) :param gpu_id: 在哪一块gpu上运行 ''' self.device = "cuda:0" checkpoint = torch.load(model_path, map_location=self.device) config = checkpoint['config'] config['arch']['backbone']['pretrained'] = False self.model = build_model(config['arch']) self.post_process = get_post_processing(config['post_processing']) self.post_process.box_thresh = post_p_thre self.img_mode = config['dataset']['train']['dataset']['args']['img_mode'] self.model.load_state_dict(checkpoint['state_dict']) self.model.to(self.device) self.model.eval() self.transform = [] for t in config['dataset']['train']['dataset']['args']['transforms']: if t['type'] in ['ToTensor', 'Normalize']: self.transform.append(t) self.transform = get_transforms(self.transform)
def __init__(self, model_path, gpu_id=None): self.gpu_id = gpu_id if self.gpu_id is not None and isinstance( self.gpu_id, int) and torch.cuda.is_available(): self.device = torch.device("cuda:%s" % self.gpu_id) else: self.device = torch.device("cpu") print('device:', self.device) checkpoint = torch.load(model_path, map_location=self.device) config = checkpoint['config'] config['arch']['backbone']['pretrained'] = False self.model = build_model(config['arch']) self.post_process = get_post_processing(config['post_processing']) self.model.load_state_dict(checkpoint['state_dict']) self.model.to(self.device) self.model.eval() self.transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) ])
def __init__(self, model_path, post_p_thre=0.7, gpu_id=None): ''' 初始化pytorch模型 :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件) :param gpu_id: 在哪一块gpu上运行 ''' self.gpu_id = gpu_id if self.gpu_id is not None and isinstance( self.gpu_id, int) and torch.cuda.is_available(): self.device = torch.device("cuda:%s" % self.gpu_id) else: self.device = torch.device("cpu") print('device:', self.device) checkpoint = torch.load(model_path, map_location=self.device) config = yaml.load( open( '/content/DBNet.pytorch/config/icdar2015_dcn_resnet18_FPN_DBhead_polyLR.yaml', 'r')) config['arch']['backbone']['pretrained'] = False self.model = build_model(config['arch']) self.post_process = get_post_processing(config['post_processing']) self.post_process.box_thresh = post_p_thre self.img_mode = config['dataset']['train']['dataset']['args'][ 'img_mode'] self.model.load_state_dict( torch.load('/content/DBNet.pytorch/ic15_resnet18')) self.model.to(self.device) self.model.eval() self.transform = [] for t in config['dataset']['train']['dataset']['args']['transforms']: if t['type'] in ['ToTensor', 'Normalize']: self.transform.append(t) self.transform = get_transforms(self.transform)
def __init__(self, model_path, post_p_thre=0.7, gpu_id=None, imageH=960, imageW=480): ''' 初始化pytorch模型, 转换tensorRT engine :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件) :param gpu_id: 在哪一块gpu上运行 ''' self.gpu_id = gpu_id if self.gpu_id is not None and isinstance( self.gpu_id, int) and torch.cuda.is_available(): self.device = torch.device("cuda:%s" % self.gpu_id) else: self.device = torch.device("cpu") print('device:', self.device) self.model_path = model_path checkpoint = torch.load(model_path, map_location=self.device) config = checkpoint['config'] config['arch']['backbone']['pretrained'] = False self.model = build_model(config['arch']) self.model.forward = self.model.forward4trt # comment F.interpolate() of 'biliner' mode config['post_processing']['args']['unclip_ratio'] = 1.8 self.post_process = get_post_processing(config['post_processing']) self.post_process.box_thresh = post_p_thre self.img_mode = config['dataset']['train']['dataset']['args'][ 'img_mode'] self.model.load_state_dict(checkpoint['state_dict']) self.model.to(self.device) self.model.eval() self.imageH = imageH self.imageW = imageW self.batch_size = 1 self.input_names = ['inputs'] self.output_names = ['pred_maps'] # ============================================================================================================= # 为了达到动态输入尺寸的目的,我们加载两个不同输入尺寸的engine, 640x640 和 960x480 # ===== 640x640 ===== input_h, input_w = (imageH + imageW) * 4 // 9, (imageH + imageW) * 4 // 9 # Convert to onnx if not os.path.exists( model_path.replace('.pth', '_{}_{}.onnx'.format( input_h, input_w))): dummy_input = torch.randn(self.batch_size, 3, input_h, input_w).to(self.device) torch.onnx.export(self.model, dummy_input, model_path.replace( '.pth', '_{}_{}.onnx'.format(input_h, input_w)), input_names=self.input_names, output_names=self.output_names, verbose=True, opset_version=9) print('Converted to onnx model, save path {}!'.format( model_path.replace('.pth', '_{}_{}.onnx'.format(input_h, input_w)))) # Convert to tensorRT engine self.engine_1x1 = get_engine( model_path.replace('.pth', '_{}_{}.onnx'.format(input_h, input_w)), model_path.replace('.pth', '_{}_{}.engine'.format(input_h, input_w))) print('Converted to tensorRT engine, save path {}!'.format( model_path.replace('.pth', '_{}_{}.engine'.format(input_h, input_w)))) self.context_1x1 = self.engine_1x1.create_execution_context() # ===== 960x480 ===== input_h, input_w = imageH, imageW # Convert to onnx if not os.path.exists( model_path.replace('.pth', '_{}_{}.onnx'.format( input_h, input_w))): dummy_input = torch.randn(self.batch_size, 3, input_h, input_w).to(self.device) torch.onnx.export(self.model, dummy_input, model_path.replace( '.pth', '_{}_{}.onnx'.format(input_h, input_w)), input_names=self.input_names, output_names=self.output_names, verbose=True, opset_version=9) print('Converted to onnx model, save path {}!'.format( model_path.replace('.pth', '_{}_{}.onnx'.format(input_h, input_w)))) # Convert to tensorRT engine self.engine_2x1 = get_engine( model_path.replace('.pth', '_{}_{}.onnx'.format(input_h, input_w)), model_path.replace('.pth', '_{}_{}.engine'.format(input_h, input_w))) print('Converted to tensorRT engine, save path {}!'.format( model_path.replace('.pth', '_{}_{}.engine'.format(input_h, input_w)))) self.context_2x1 = self.engine_2x1.create_execution_context() # ============================================================================================================= self.transform = [] for t in config['dataset']['train']['dataset']['args']['transforms']: if t['type'] in ['ToTensor', 'Normalize']: self.transform.append(t) self.transform = get_transforms(self.transform)
def train(): config = Config() train_config = TrainConfig() params = init_args() device = "cuda" if torch.cuda.is_available() else "cpu" # 检查设备 # 加载训练数据集 train_kwargs = { 'num_workers': 4, 'pin_memory': True } if torch.cuda.is_available() else {} train_dataset = ICDAR2015Dataset(params=params, mode="train", img_transform=transforms.Compose([ ImageAug(), EastRandomCropData(config), MakeBorderMap(config), MakeShrinkMap(config) ]), transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])) training_data_batch = DataLoader(train_dataset, batch_size=config.train_batch_size, shuffle=True, drop_last=True, **train_kwargs) # 加载测试数据集 icdar_collect_fn = ICDARCollectFN() # 自定义Batch生成类 test_kwargs = { 'num_workers': 1, 'pin_memory': False, "collate_fn": icdar_collect_fn } if torch.cuda.is_available() else {} test_dataset = ICDAR2015Dataset( params=params, mode="test", img_transform=transforms.Compose([ResizeShortSize(config)]), transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])) testing_data_batch = DataLoader(test_dataset, batch_size=config.test_batch_size, shuffle=False, drop_last=False, **test_kwargs) model = DBNet(params=params) criterion = DBLoss(params=params).to(device) # 生成配置文件 post_processing_arg_dict = { "thresh": train_config.post_processing_thresh, "box_thresh": train_config.post_processing_box_thresh, "max_candidates": train_config.post_processing_max_candidates, "unclip_ratio": train_config.post_processing_unclip_ratio # from paper } post_processing_arg = { "args": post_processing_arg_dict, "type": train_config.post_processing_type } metric_arg_dict = { "is_output_polygon": train_config.metric_is_output_polygon } metric_arg = {"type": train_config.metric_type, "args": metric_arg_dict} val_config = {"post_processing": post_processing_arg, "metric": metric_arg} post_p = get_post_processing(val_config['post_processing']) metric = get_metric(val_config['metric']) train_helper = Train_Helper(train_config=train_config, model=model, criterion=criterion, train_loader=training_data_batch, validate_loader=testing_data_batch, metric_cls=metric, post_process=post_p) train_helper.train()
def main(config): import torch from models import build_model, build_loss from data_loader import get_dataloader from trainer import Trainer from post_processing import get_post_processing from utils import get_metric from utils import setup_logger if torch.cuda.device_count() > 1: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group( backend="nccl", init_method="env://", world_size=torch.cuda.device_count(), rank=args.local_rank) config['distributed'] = True else: config['distributed'] = False config['local_rank'] = args.local_rank config['arch']['backbone']['in_channels'] = 3 if config['dataset'][ 'train']['dataset']['args']['img_mode'] != 'GRAY' else 1 model = build_model(config['arch']) if config['local_rank'] == 0: save_dir = os.path.join(config['trainer']['output_dir'], config['name'] + '_' + model.name) if not os.path.exists(save_dir): os.makedirs(save_dir) logger = setup_logger(os.path.join(save_dir, 'train.log')) if 'evolve' in config.keys( ) and config['evolve']['flag'] and not config['distributed']: meta = { 'optimizer.args.lr': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3) 'lr_scheduler.args.warmup_epoch': (1, 0, 5), 'loss.alpha': (1, 0.5, 3), 'loss.beta': (2, 5, 20), 'loss.ohem_ratio': (1, 1, 5), 'post_processing.args.box_thresh': (0.3, 0.4, 1.0), 'dataset.train.dataset.args.pre_processes.[1].args.min_crop_side_ratio': (1, 0.1, 0.9), 'dataset.train.dataset.args.pre_processes.[2].args.thresh_max': (0.3, 0.4, 1.0), } # image mixup (probability) config['notest'] = True config['nosave'] = True saved_path = os.path.join(config['trainer']['output_dir'], config['name'] + '_' + model.name) if not os.path.exists(os.path.join(saved_path, 'evolve')): os.makedirs(os.path.join(saved_path, 'evolve')) yaml_file = os.path.join(saved_path, 'evolve', 'hyp_evolved.yaml') evolve_file = os.path.join(saved_path, 'evolve', 'evolve.txt') for _ in range(300): if os.path.exists(evolve_file): parent = 'single' x = np.loadtxt(evolve_file, ndmin=2) n = min(5, len(x)) x = x[np.argsort(-fitness(x))][:n] w = fitness(x) - fitness(x).min() if len(x) == 1: x = x[0] elif parent == 'single': # x = x[random.randint(0, n - 1)] # random selection x = x[random.choices(range(n), weights=w)[0]] # weighted selection elif parent == 'weighted': x = (x * w.reshape( n, 1)).sum(0) / w.sum() # weighted combination # Mutate mp, s = 0.8, 0.2 # mutation probability, sigma npr = np.random npr.seed(int(time.time())) g = np.array([x[0] for x in meta.values()]) # gains 0-1 ng = len(meta) v = np.ones(ng) while all( v == 1 ): # mutate until a change occurs (prevent duplicates) v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0) # for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300) # hyp[k] = float(x[i + 7] * v[i]) # mutate for i, k in enumerate(meta.keys()): config_keys = k.split('.') str_config = 'config' for config_key in config_keys: if config_key.startswith('[') and config_key.endswith( ']'): str_config = str_config + config_key else: str_config = str_config + '[\'' + config_key + '\']' exec(str_config + '=x[i]*v[i]') meta_value = [] for k, v in meta.items(): config_keys = k.split('.') str_config = 'config' for config_key in config_keys: if config_key.startswith('[') and config_key.endswith(']'): str_config = str_config + config_key else: str_config = str_config + '[\'' + config_key + '\']' # str_config = 'config[\'' + '\'][\''.join(k.split('.')) + '\']' exec('print(' + str_config + ')') exec(str_config + '=max(' + str_config + ', v[1])') exec(str_config + ' = min(' + str_config + ', v[2])') exec(str_config + ' = round(' + str_config + ', 5)') exec('meta_value.append(' + str_config + ')') # hyp[k] = max(hyp[k], v[1]) # lower limit # hyp[k] = min(hyp[k], v[2]) # upper limit # hyp[k] = round(hyp[k], 5) # significant digits train_loader = get_dataloader(config['dataset']['train'], config['distributed']) assert train_loader is not None if 'validate' in config['dataset']: validate_loader = get_dataloader(config['dataset']['validate'], False) else: validate_loader = None criterion = build_loss(config['loss']).cuda() post_p = get_post_processing(config['post_processing']) metric = get_metric(config['metric']) trainer = Trainer( config=config, model=model, criterion=criterion, train_loader=train_loader, post_process=post_p, metric_cls=metric, validate_loader=validate_loader, logger=(logger if config['local_rank'] == 0 else None)) results = trainer.train() print_mutation(results, yaml_file, evolve_file, meta_value) else: train_loader = get_dataloader(config['dataset']['train'], config['distributed']) assert train_loader is not None if 'validate' in config['dataset']: validate_loader = get_dataloader(config['dataset']['validate'], False) else: validate_loader = None criterion = build_loss(config['loss']).cuda() post_p = get_post_processing(config['post_processing']) metric = get_metric(config['metric']) trainer = Trainer( config=config, model=model, criterion=criterion, train_loader=train_loader, post_process=post_p, metric_cls=metric, validate_loader=validate_loader, logger=(logger if config['local_rank'] == 0 else None)) trainer.train()
def __init__(self, model_path, post_p_thre=0.7, gpu_id=None, short_size=640): ''' 初始化pytorch模型, 转换tensorRT engine :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件) :param gpu_id: 在哪一块gpu上运行 ''' self.gpu_id = gpu_id if self.gpu_id is not None and isinstance( self.gpu_id, int) and torch.cuda.is_available(): self.device = torch.device("cuda:%s" % self.gpu_id) else: self.device = torch.device("cpu") print('device:', self.device) self.model_path = model_path checkpoint = torch.load(model_path, map_location=self.device) config = checkpoint['config'] config['arch']['backbone']['pretrained'] = False self.model = build_model(config['arch']) self.model.forward = self.model.forward4trt # comment F.interpolate() of 'biliner' mode config['post_processing']['args']['unclip_ratio'] = 1.8 self.post_process = get_post_processing(config['post_processing']) self.post_process.box_thresh = post_p_thre self.img_mode = config['dataset']['train']['dataset']['args'][ 'img_mode'] self.model.load_state_dict(checkpoint['state_dict']) self.model.to(self.device) self.model.eval() self.short_size = short_size self.batch_size = 1 self.input_names = ['inputs'] self.output_names = ['pred_maps'] # ============================================================================================================= # Convert to onnx if not os.path.exists(model_path.replace('.pth', '_dynamic.onnx')): dummy_input = torch.randn(self.batch_size, 3, self.short_size, self.short_size).to(self.device) # dynamic_axes = {self.input_names[0]: {2:'width', 3:'height'}, # self.output_names[0]: {2:'width', 3:'height'}} torch.onnx.export( self.model, dummy_input, model_path.replace('.pth', '_dynamic.onnx'), # dynamic_axes= dynamic_axes, input_names=self.input_names, output_names=self.output_names, verbose=True, opset_version=10) # onnx_model = onnx.load(model_path.replace('.pth', '_dynamic.onnx')) # onnx.checker.check_model(onnx_model) print('Converted to onnx model, save path {}!'.format( model_path.replace('.pth', '_dynamic.onnx'))) # Convert to tensorRT engine self.engine = get_engine(model_path.replace('.pth', '_dynamic.onnx'), model_path.replace('.pth', '_dynamic.engine')) print('Converted to tensorRT engine, save path {}!'.format( model_path.replace('.pth', '_dynamic.engine'))) self.context = self.engine.create_execution_context() # ============================================================================================================= self.transform = [] for t in config['dataset']['train']['dataset']['args']['transforms']: if t['type'] in ['ToTensor', 'Normalize']: self.transform.append(t) self.transform = get_transforms(self.transform)
def eval(): config = Config() eval_config = EvalConfig() params = init_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 检查设备 # 加载测试数据集 icdar_collect_fn = ICDARCollectFN() # 自定义Batch生成类 test_kwargs = {'num_workers': 1, 'pin_memory': False, "collate_fn": icdar_collect_fn} if torch.cuda.is_available() else {} test_dataset = ICDAR2015Dataset(params=params, mode="test", img_transform=transforms.Compose( [ResizeShortSize(config)]), transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])) testing_data_batch = DataLoader(test_dataset, batch_size=config.test_batch_size, shuffle=False, drop_last=False, **test_kwargs) checkpoint = torch.load(eval_config.model_path, map_location=torch.device('cpu')) # model = DBNet(params=params) # model.load_state_dict(checkpoint['state_dict']) # model.to(device) model = torch.load(eval_config.model_path_model) model.to(device) # 生成配置文件 post_processing_arg_dict = { "thresh": eval_config.post_processing_thresh, "box_thresh": eval_config.post_processing_box_thresh, "max_candidates": eval_config.post_processing_max_candidates, "unclip_ratio": eval_config.post_processing_unclip_ratio # from paper } post_processing_arg = { "args": post_processing_arg_dict, "type": eval_config.post_processing_type } metric_arg_dict = { "is_output_polygon": eval_config.metric_is_output_polygon } metric_arg = { "type": eval_config.metric_type, "args": metric_arg_dict } val_config = { "post_processing": post_processing_arg, "metric": metric_arg } post_process = get_post_processing(val_config['post_processing']) metric_cls = get_metric(val_config['metric']) model.eval() # torch.cuda.empty_cache() # speed up evaluating after training finished raw_metrics = [] total_frame = 0.0 total_time = 0.0 for i, batch in tqdm(enumerate(testing_data_batch), total=len(testing_data_batch), desc='test model'): with torch.no_grad(): # 数据进行转换和丢到gpu for key, value in batch.items(): if value is not None: if isinstance(value, torch.Tensor): batch[key] = value.to(device) start = time.time() preds = model(batch['img']) boxes, scores = post_process(batch, preds, is_output_polygon=metric_cls.is_output_polygon) total_frame += batch['img'].size()[0] total_time += time.time() - start raw_metric = metric_cls.validate_measure(batch, (boxes, scores)) raw_metrics.append(raw_metric) metrics = metric_cls.gather_measure(raw_metrics) print('FPS:{}'.format(total_frame / total_time)) return metrics['recall'].avg, metrics['precision'].avg, metrics['fmeasure'].avg