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()