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
