def train_batch(self,
inputs,
targets,
weights=None,
update=True,
logname="train"):
if update:
self.set_training(True)
self.optimizer.zero_grad()
else:
self.set_training(False)
self.set_inputs(inputs)
self.forward()
if weights is not None:
self.cuweights = autograd.Variable(torch.randn(1, 1).cuda())
dlh.assign(self.cuweights, weights, False)
self.cuoutput = self.weighted(self.cuoutput, self.cuweights)
culoss = self.compute_loss(targets, weights=weights)
if update:
culoss.backward()
self.optimizer.step()
ploss = dlh.novar(culoss)[0]
self.ntrain += dlh.size(inputs, 0)
add_log(self.log,
logname,
loss=ploss,
ntrain=self.ntrain,
lr=self.current_lr)
return self.get_outputs(), ploss
def ctc_align(prob, target):
"""Perform CTC alignment on torch sequence batches (using ocrolstm).
Inputs are in BDL format.
"""
import cctc
assert dlh.sequence_is_normalized(prob), prob
assert dlh.sequence_is_normalized(target), target
# inputs are BDL
prob_ = dlh.novar(prob).permute(0, 2, 1).cpu().contiguous()
target_ = dlh.novar(target).permute(0, 2, 1).cpu().contiguous()
# prob_ and target_ are both BLD now
assert prob_.size(0) == target_.size(0), (prob_.size(), target_.size())
assert prob_.size(2) == target_.size(2), (prob_.size(), target_.size())
assert prob_.size(1) >= target_.size(1), (prob_.size(), target_.size())
result = torch.rand(1)
cctc.ctc_align_targets_batch(result, prob_, target_)
return dlh.typeas(result.permute(0, 2, 1).contiguous(), prob)
def get_outputs(self):
"""Performs any necessary transformations on the output tensor.
"""
return dlh.novar(self.cuoutput).cpu()