Method forward_backward

plot-pp/trainers/plotpp.py:333–355

(self, batch)

        self.scaler = GradScaler() if cfg.TRAINER.PLOTPP.PREC == "amp" else None
        
    def forward_backward(self, batch):
        image, label = self.parse_batch_train(batch)
        
        prec = self.cfg.TRAINER.PLOTPP.PREC
        if prec == "amp":
            with autocast():
                output = self.model(image, label)
            self.optim.zero_grad()
            self.scaler.scale(output).backward()
            self.scaler.step(self.optim)
            self.scaler.update()
        else:
            output = self.model(image)
            loss = F.cross_entropy(output, label)
            self.model_backward_and_update(loss)
            
        loss_summary = {"loss": loss.item(),
                         "acc": compute_accuracy(output, label)[0].item()}
        
        if (self.batch_idx + 1) == self.num_batches:
            self.update_lr()
            
        return loss_summary
    
    def parse_batch_train(self, batch):
        input = batch["img"]