def train_cnocr(args):
head = '%(asctime)-15s %(message)s'
logging.basicConfig(level=logging.DEBUG, format=head)
args.model_name = args.emb_model_type + '-' + args.seq_model_type
out_dir = os.path.join(args.out_model_dir, args.model_name)
logger.info('save models to dir: %s' % out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
args.prefix = os.path.join(
out_dir, 'cnocr-v{}-{}'.format(__version__, args.model_name))
hp = CnHyperparams()
hp = _update_hp(hp, args)
network, hp = gen_network(args.model_name, hp)
metrics = CtcMetrics(hp.seq_length)
data_train, data_val = _gen_iters(hp, args.train_file, args.test_file,
args.use_train_image_aug, args.dataset,
args.charset, args.debug)
data_names = ['data']
fit(
network=network,
data_train=data_train,
data_val=data_val,
metrics=metrics,
args=args,
hp=hp,
data_names=data_names,
)
def test_gen_networks(model_name):
logger.info('model_name: %s', model_name)
network, hp = gen_network(model_name, HP)
shape_dict = get_infer_shape(network, HP)
logger.info('shape_dict: %s', shape_dict)
assert shape_dict['pred_fc_output'] == (
hp.batch_size * hp.seq_length,
hp.num_classes,
)
def _get_module(self):
network, self._hp = gen_network(self._model_name, self._hp)
hp = self._hp
prefix = os.path.join(self._model_dir, self._model_file_prefix)
data_names = ['data']
data_shapes = [(data_names[0], (hp.batch_size, 1, hp.img_height,
hp.img_width))]
logger.info('loading model parameters from dir %s' % self._model_dir)
mod = load_module(prefix,
self._model_epoch,
data_names,
data_shapes,
network=network)
return mod