converter = utils.strLabelConverter(opt.alphabet)
criterion = torch.nn.CTCLoss()
# custom weights initialization called on crnn
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
crnn = crnn.CRNN(opt.imgH, nc, nclass, opt.nh)
crnn.apply(weights_init)
if opt.pretrained != '':
print('loading pretrained model from %s' % opt.pretrained)
crnn.load_state_dict(torch.load(opt.pretrained))
print(crnn)
image = torch.FloatTensor(opt.batchSize, 3, opt.imgH, opt.imgH)
text = torch.IntTensor(opt.batchSize * 5)
length = torch.IntTensor(opt.batchSize)
if opt.cuda:
crnn.cuda()
crnn = torch.nn.DataParallel(crnn, device_ids=range(opt.ngpu))
image = image.cuda()
criterion = criterion.cuda()
def train_net(args):
manual_seed = 7
random.seed(manual_seed)
np.random.seed(manual_seed)
torch.manual_seed(manual_seed)
checkpoint = args.checkpoint
start_epoch = 0
best_loss = float('inf')
writer = SummaryWriter()
epochs_since_improvement = 0
# custom weights initialization called on crnn
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
# Initialize / load checkpoint
if checkpoint is None:
model = CRNN(imgH, num_channels, num_classes, num_hidden)
model.apply(weights_init)
# model = nn.DataParallel(model)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
betas=(0.9, 0.98),
eps=1e-09)
# optimizer = CRNNOptimizer(
# torch.optim.Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-09),
# args.k,
# num_hidden,
# args.warmup_steps)
else:
checkpoint = torch.load(checkpoint)
start_epoch = checkpoint['epoch'] + 1
epochs_since_improvement = checkpoint['epochs_since_improvement']
model = checkpoint['model']
optimizer = checkpoint['optimizer']
logger = utils.get_logger()
# Move to GPU, if available
model = model.to(device)
# Loss function
criterion = nn.CTCLoss(reduction='mean').to(device)
# Custom dataloaders
train_dataset = data_gen.MJSynthDataset('train')
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=num_workers)
valid_dataset = data_gen.MJSynthDataset('val')
valid_loader = torch.utils.data.DataLoader(valid_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=num_workers)
# Epochs
for epoch in range(start_epoch, args.end_epoch):
# One epoch's training
train_loss, train_acc = train(train_loader=train_loader,
model=model,
criterion=criterion,
optimizer=optimizer,
epoch=epoch,
logger=logger)
writer.add_scalar('Train_Loss', train_loss, epoch)
writer.add_scalar('Train_Accuracy', train_acc, epoch)
# One epoch's validation
valid_loss, valid_acc = valid(valid_loader=valid_loader,
model=model,
criterion=criterion,
logger=logger)
writer.add_scalar('Validation_Loss', valid_loss, epoch)
writer.add_scalar('Validation_Accuracy', valid_acc, epoch)
# Check if there was an improvement
is_best = valid_loss < best_loss
best_loss = min(valid_loss, best_loss)
if not is_best:
epochs_since_improvement += 1
print("\nEpochs since last improvement: %d\n" %
(epochs_since_improvement, ))
else:
epochs_since_improvement = 0
# Save checkpoint
utils.save_checkpoint(epoch, epochs_since_improvement, model,
optimizer, best_loss, is_best)
def train_net(args):
manual_seed = 7
random.seed(manual_seed)
np.random.seed(manual_seed)
torch.manual_seed(manual_seed)
checkpoint = args.checkpoint
start_epoch = 0
best_loss = float('inf')
writer = SummaryWriter()
epochs_since_improvement = 0
# custom weights initialization called on crnn
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
# Initialize / load checkpoint
if checkpoint is None:
model = CRNN(imgH, nc, nclass, nh)
model.apply(weights_init)
# model = nn.DataParallel(model)
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.mom,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
else:
checkpoint = torch.load(checkpoint)
start_epoch = checkpoint['epoch'] + 1
epochs_since_improvement = checkpoint['epochs_since_improvement']
model = checkpoint['model']
optimizer = checkpoint['optimizer']
logger = utils.get_logger()
# Move to GPU, if available
model = model.to(device)
# Loss function
criterion = nn.CTCLoss(reduction='mean').to(device)
# Custom dataloaders
train_dataset = data_gen.Ic2015Dataset('train')
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=num_workers)
test_dataset = data_gen.Ic2015Dataset('test')
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=num_workers)
# Epochs
for epoch in range(start_epoch, args.end_epoch):
# One epoch's training
train_loss, train_acc = train(train_loader=train_loader,
model=model,
criterion=criterion,
optimizer=optimizer,
epoch=epoch,
logger=logger)
effective_lr = utils.get_learning_rate(optimizer)
print('\nCurrent effective learning rate: {}\n'.format(effective_lr))
writer.add_scalar('Learning_Rate', effective_lr, epoch)
writer.add_scalar('Train_Loss', train_loss, epoch)
writer.add_scalar('Train_Accuracy', train_acc, epoch)
# One epoch's validation
test_loss, test_acc = test(test_loader=test_loader,
model=model,
criterion=criterion,
logger=logger)
writer.add_scalar('Test_Loss', test_loss, epoch)
writer.add_scalar('Test_Accuracy', test_acc, epoch)
# Check if there was an improvement
is_best = test_loss < best_loss
best_loss = min(test_loss, best_loss)
if not is_best:
epochs_since_improvement += 1
print("\nEpochs since last improvement: %d\n" %
(epochs_since_improvement, ))
else:
epochs_since_improvement = 0
# Save checkpoint
utils.save_checkpoint(epoch, epochs_since_improvement, model,
optimizer, best_loss, is_best)
