def training_step()

in src/sagemaker_defect_detection/detector.py [0:0]

• 19 lines of code
• 11 McCabe index (conditional complexity)

    def training_step(self, batch, batch_idx):
        images, targets, _ = batch
        if self.train_rpn:
            targets = [{"boxes": t["boxes"]} for t in targets]
            _, loss_dict = self(images, targets=targets)
            loss = sum(loss for loss in loss_dict.values())
            return OrderedDict({"loss": loss, "progress_bar": loss_dict, "log": loss_dict})

        elif self.train_roi:
            _, loss_dict = self(images, targets=targets)
            loss = sum(loss for loss in loss_dict.values())
            return OrderedDict({"loss": loss, "progress_bar": loss_dict, "log": loss_dict})
        else:
            images = list(image for image in images)
            targets = [{k: v for k, v in t.items()} for t in targets]
            loss_dict = self(images, targets=targets)
            # loss keys: ['loss_classifier', 'loss_box_reg', 'loss_objectness', 'loss_rpn_box_reg']
            loss = sum(loss for loss in loss_dict.values())
            if not math.isfinite(loss.item()):
                sys.exit(1)

            return OrderedDict({"loss": loss, "progress_bar": loss_dict, "log": loss_dict})