def __init__(self, data_dir, key, batch_size=32, debug=False) -> None:
"""
IAM-OnDB Handwritten character recognition.
A many-to-one sequence problem, typically solved with RNN+CTC loss.
refs: https://papers.nips.cc/paper/3213-unconstrained-on-line-handwriting-recognition-with-recurrent-neural-networks
"""
super().__init__()
self.key = key
self.batch_size = batch_size
if not path.exists(data_dir + f'/x_{key}.npy'):
print("Loading from raw data.")
f_ids = data_dir + f'/ids_{key}.txt'
file_names = [f.strip() for f in open(f_ids, mode='r').readlines()]
self.xs, self.ys = fetch_iamondb_from_list(data_dir, file_names)
np.save(data_dir + f'/x_{key}.npy', self.xs, allow_pickle=True)
np.save(data_dir + f'/y_{key}.npy', self.ys, allow_pickle=True)
else:
print("Loading numpy arrays.")
self.xs = np.load(data_dir + f'/x_{key}.npy', allow_pickle=True)
self.ys = np.load(data_dir + f'/y_{key}.npy', allow_pickle=True)
self.xs = [cuda_move(torch.tensor(el)) for el in self.xs]
self.ys = [cuda_move(torch.tensor(el)) for el in self.ys]
if debug:
self.xs = self.xs[:200]
self.ys = self.ys[:200]
def after_train_before_validate(self, model_trainer):
model_name = self.log_dir + 'best_model'
if os.path.isfile(model_name + '.pt'):
# model_trainer.model = cuda_move(torch.load(model_name + '.pt'))
model_trainer.model.load_state_dict(torch.load(model_name + '.pt'))
model_trainer.logger.info("Loaded best model checkpoint before final validation.")
elif os.path.isfile(model_name + '.ptj'):
try:
model_trainer.model = cuda_move(torch.jit.load(model_name + '.ptj'))
model_trainer.logger.info("Loaded best model checkpoint before final validation.")
except BaseException as e:
model_trainer.logger.info(str(e))
model_trainer.logger.info("Could not load model. Checkpoint is corrupted.")
else:
model_trainer.logger.info("Checkpoint file not found. Using current model for final validation.")
def iter(self):
n_samples = len(self.xs)
if self.key == 'train': # shuffle dataset
idxs = torch.randperm(n_samples)
self.xs = [self.xs[ii] for ii in idxs]
self.ys = [self.ys[ii] for ii in idxs]
for ii in range(0, len(self.xs), self.batch_size):
x_batch = self.xs[ii:ii + self.batch_size]
y_batch = self.ys[ii:ii + self.batch_size]
bs = len(x_batch)
x_t_max = max(el.shape[0] for el in x_batch)
x = torch.zeros(x_t_max, bs, x_batch[0].shape[-1])
y_t_max = max(el.shape[0] for el in y_batch)
y = torch.zeros(bs, y_t_max)
for bi, (x_el, y_el) in enumerate(zip(x_batch, y_batch)):
x[:x_el.shape[0], bi] = x_el
y[bi, :y_el.shape[0]] = y_el
x = cuda_move(x)
y = cuda_move(y).int()
x_t_batch = cuda_move(
torch.tensor([el.shape[0] for el in x_batch]).int())
y_t_batch = cuda_move(
torch.tensor([el.shape[0] for el in y_batch]).int())
yield x, (y, x_t_batch, y_t_batch)
def __str__(self):
return f"IAMOnDB Dataset (key={self.key}, bs={self.batch_size})"
if __name__ == '__main__':
data_dir = '/home/USER/data/iamondb'
# Load model and compute WER
valid_data = IAMOnDB(data_dir, key='valid')
model = ItemClassifier(rnn=LSTMLayer(4, 48),
hidden_size=48 * 9,
output_size=58)
model.load_state_dict(
torch.load('./logs/iamondb/lstm_gs0/k_0/best_model.pt'))
model = cuda_move(model)
acc = 0
bi = 0
for batch in tqdm(valid_data.iter()):
y_pred = model(batch[0])
acc += valid_data.metric_score(batch, y_pred)
bi += 1
acc = acc / bi
print(f"ACC: {acc}")
exit()
# Dataset statistics
data = fetch_iamondb(data_dir)
(train_x, train_y), (valid_x, valid_y), (test_x, test_y) = data
train_x = np.vstack(train_x)
def train_foo(log_dir, params):
DEBUG = params['DEBUG']
model_type = params['model_type']
train_data = IAMOnDB(data_dir, 'train', params['batch_size'], debug=DEBUG)
val_data = IAMOnDB(data_dir, 'valid', params['batch_size'], debug=DEBUG)
test_data = IAMOnDB(data_dir, 'test', params['batch_size'], debug=DEBUG)
if DEBUG:
params['n_epochs'] = 15
params['pretrain_every'] = 5
params['k_unroll'] = 3
params['ulm_n_epochs'] = 4
ins, hs, ms = params['in_size'], params['hidden_size'], params['memory_size']
if model_type == 'lstm':
rnn = BidirectionalRNN(LSTMLayer(ins, hs), LSTMLayer(ins, hs))
ro_in = 2 * hs
cbs, regs = [], []
elif model_type == 'lmn':
rnn = BidirectionalRNN(LMNLayer(ins, hs, ms), LMNLayer(ins, hs, ms))
ro_in = 2 * ms
cbs = [OrthogonalInit([rnn.fw_layer.Wmm, rnn.bw_layer.Wmm])]
regs = [OrthogonalPenalty(params['ortho_penalty'], [rnn.fw_layer.Wmm, rnn.bw_layer.Wmm])]
elif model_type == 'ms_lmn':
rnn = BidirectionalRNN(
MultiScaleLMN(ins, hs, ms, num_modules=params['num_modules']),
MultiScaleLMN(ins, hs, ms, num_modules=params['num_modules'])
)
ro_in = 2 * params['num_modules'] * ms
cbs, regs = [], []
elif model_type == 'cw_rnn':
rnn = BidirectionalRNN(
ClockworkRNN(ins, hs, params['num_modules']),
ClockworkRNN(ins, hs, params['num_modules'])
)
ro_in = 2 * hs * params['num_modules']
cbs, regs = [], []
elif model_type == 'pret_ms_lmn':
rnn = BidirectionalRNN(
MultiScaleLMN(ins, hs, ms, num_modules=0, max_modules=params['num_modules']),
MultiScaleLMN(ins, hs, ms, num_modules=0, max_modules=params['num_modules'])
)
ro_in = 2 * ms * params['num_modules']
cbs = [IncrementalTrainingCallback(params, [rnn.bw_layer, rnn.fw_layer], train_data, val_data, test_data, log_dir)]
regs = []
model = cuda_move(ItemClassifier(
rnn=rnn,
hidden_size=ro_in,
output_size=params['output_size']
))
if params['optimizer'] == 'adam':
optimizer = torch.optim.Adam(model.parameters(),
lr=params['learning_rate'],
weight_decay=params['l2_decay'])
trainer = SequentialTaskTrainer(
cuda_move(model),
n_epochs=params['n_epochs'],
optimizer=optimizer,
regularizers=regs,
log_dir=log_dir,
callbacks=cbs,
use_tqdm=True,
validation_steps=1,
patience=params['patience'],
checkpoint_mode='loss',
debug=DEBUG,
)
trainer.append_hyperparam_dict(params)
trainer.fit(train_data, val_data)
e, a = trainer.compute_metrics(test_data)
trainer.logger.info(f"Final performance on TEST set: loss {e}, acc {a}")
trainer.logger.info(trainer.best_result)
return trainer.best_result