cudnn.benchmark = True
if torch.cuda.is_available() and not opt.cuda:
print(
"WARNING: You have a CUDA device, so you should probably run with --cuda"
)
transform = None
if opt.trainlist.lower().endswith(".h5"):
train_dataset = h5dataset.H5Dataset(opt.trainlist)
opt.imgW, h = train_dataset.imaDims()
opt.imgH = int(16 * round(h / 16))
print(" ** TrainSize:", (opt.imgW, opt.imgH))
opt.max_width = train_dataset.prof_lenght
else:
train_dataset = dataset.listDataset(list_file=opt.trainlist,
transform=transform)
print(" -- Saving parameters json:")
with open(opt.outdir + os.sep + "params.json", "w") as p:
p.write(json.dumps(vars(opt)))
assert train_dataset
if not opt.random_sample:
sampler = dataset.randomSequentialSampler(train_dataset, opt.batchSize)
else:
sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.batchSize,
shuffle=False,
sampler=sampler,
num_workers=int(opt.workers),
SOS_TOKEN = 0
EOS_TOKEN = 1 # end of sequence
BLANK = 2 # blank character
opt.manualSeed = random.randint(1, 10000) # fix seed
random.seed(opt.manualSeed)
np.random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
cudnn.benchmark = True
if opt.cuda:
torch.cuda.set_device(opt.gpuid)
transform = None
train_dataset = dataset.listDataset(list_file=opt.trainList,
transform=transform)
assert train_dataset
if not opt.random_sample:
sampler = dataset.randomSequentialSampler(train_dataset, opt.batchSize)
else:
sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.batchSize,
shuffle=False,
sampler=sampler,
num_workers=int(opt.workers),
collate_fn=dataset.alignCollate(
height=opt.height,
width=opt.width,
keep_ratio=opt.keep_ratio))
target_variable[1:, count]):
if pred == target:
n_correct += 1
texts = cpu_texts[count]
print('%d Pred:%-20s, GT: %-20s' %
(n_current, decoded_words[count], texts))
n_current += 1
accuracy = n_correct / float(n_total)
print('Loss: %f, Accuracy: %f' % (loss_avg.val(), accuracy))
if __name__ == '__main__':
test_dataset = dataset.listDataset(list_file=opt.testList,
transform=dataset.resizeNormalize(
(opt.width, opt.height)))
nclass = len(alphabet) + 3
nc = 1
criterion = torch.nn.NLLLoss()
encoder = model.encoder(opt.height, nc=nc, nh=256)
decoder = model.decoder(nh=256, nclass=nclass, dropout_p=0.1)
if opt.encoder:
print('loading pretrained encoder model from %s' % opt.encoder)
encoder.load_state_dict(torch.load(opt.encoder))
if opt.decoder:
print('loading pretrained decoder model from %s' % opt.decoder)
decoder.load_state_dict(torch.load(opt.decoder))
if opt.loadModelEpoch > 0:
encoder_path = 'model/encoder_%d.pth' % opt.loadModelEpoch