def main(config, resume):
train_logger = Logger()
if config['data_loader']['dataset'] == 'icdar2015':
# ICDAR 2015
data_root = pathlib.Path(config['data_loader']['data_dir'])
ICDARDataset2015 = ICDAR(data_root, year='2015')
data_loader = OCRDataLoaderFactory(config, ICDARDataset2015)
train = data_loader.train()
val = data_loader.val()
elif config['data_loader']['dataset'] == 'synth800k':
data_loader = SynthTextDataLoaderFactory(config)
train = data_loader.train()
val = data_loader.val()
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join([str(i) for i in config['gpus']])
model = eval(config['arch'])(config['model'])
model.summary()
loss = eval(config['loss'])(config['model'])
metrics = [eval(metric) for metric in config['metrics']]
trainer = Trainer(model, loss, metrics,
resume=resume,
config=config,
data_loader=train,
valid_data_loader=val,
train_logger=train_logger,
toolbox = Toolbox)
trainer.train()
def main(config, resume):
train_logger = Logger()
if config['data_loader']['dataset'] == 'icdar2015':
# ICDAR 2015
data_root = pathlib.Path(config['data_loader']['data_dir'])
ICDARDataset2015 = ICDAR(data_root, year='2015')
data_loader = OCRDataLoaderFactory(config, ICDARDataset2015)
train = data_loader.train()
val = data_loader.val()
elif config['data_loader']['dataset'] == 'synth800k':
data_loader = SynthTextDataLoaderFactory(config)
train = data_loader.train()
val = data_loader.val()
elif config['data_loader']['dataset'] == 'mydataset':
image_root = pathlib.Path(config['data_loader']['image_dir'])
annotation_dir = pathlib.Path(config['data_loader']['annotation_dir'])
print('alpha1. image_root = ' + str(image_root) +
', annotation_dir = ' + str(annotation_dir))
data_loader = OCRDataLoaderFactory(
config, MyDataset(image_root, annotation_dir))
train = data_loader.train()
val = data_loader.val()
#print(type(train))
#print(train)
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
[str(i) for i in config['gpus']])
model = eval(config['arch'])(config)
#model.summary()
loss = eval(config['loss'])(config['model'])
metrics = [eval(metric) for metric in config['metrics']]
finetune_model = config['finetune']
trainer = Trainer(model,
loss,
metrics,
finetune=finetune_model,
resume=resume,
config=config,
data_loader=train,
valid_data_loader=val,
train_logger=train_logger,
toolbox=Toolbox,
keys=getattr(common_str, config['model']['keys']))
print('alpha3: just before trainer.train()')
trainer.train()
print('alpha4: just after trainer.train()')
def main(config, resume):
train_logger = Logger()
if config['data_loader']['dataset'] == 'icdar2015':
# ICDAR 2015
data_root = pathlib.Path(config['data_loader']['data_dir'])
ICDARDataset2015 = ICDAR(data_root, year='2015')
data_loader = OCRDataLoaderFactory(config, ICDARDataset2015)
train = data_loader.train()
val = data_loader.val()
elif config['data_loader']['dataset'] == 'synth800k':
data_loader = SynthTextDataLoaderFactory(config)
train = data_loader.train()
val = data_loader.val()
elif config['data_loader']['dataset'] == 'mydataset':
data_root = pathlib.Path(config['data_loader']['data_dir'])
data_loader = OCRDataLoaderFactory(config, MyDataset(data_root))
train = data_loader.train()
val = data_loader.val()
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
[str(i) for i in config['gpus']])
model = eval(config['arch'])(config,
is_train=True) # 将模型的字符串转化为表达式调用，然后实例化
#model.summary() # 模型结构较为复杂，可以选择不打印模型结构，注释掉即可
loss = eval(config['loss'])(config['model'])
metrics = [eval(metric) for metric in config['metrics']]
finetune_model = config['finetune']
trainer = Trainer(
model,
loss,
metrics,
finetune=finetune_model, # 导入预训练的模型，在此基础上继续训练
resume=resume,
config=config,
data_loader=train,
valid_data_loader=val,
train_logger=train_logger,
toolbox=Toolbox,
keys=getattr(common_str, config['model']['keys']))
trainer.train()