YOLOX model module. The module list is defined by create_yolov3_modules function. The network returns loss values from three YOLO layers during training and detection results during test.
| 9 | |
| 10 | |
| 11 | class YOLOX(nn.Module): |
| 12 | """ |
| 13 | YOLOX model module. The module list is defined by create_yolov3_modules function. |
| 14 | The network returns loss values from three YOLO layers during training |
| 15 | and detection results during test. |
| 16 | """ |
| 17 | |
| 18 | def __init__(self, backbone=None, head=None): |
| 19 | super().__init__() |
| 20 | if backbone is None: |
| 21 | backbone = YOLOPAFPN() |
| 22 | if head is None: |
| 23 | head = YOLOXHead(80) |
| 24 | |
| 25 | self.backbone = backbone |
| 26 | self.head = head |
| 27 | |
| 28 | def forward(self, x, targets=None): |
| 29 | # fpn output content features of [dark3, dark4, dark5] |
| 30 | fpn_outs = self.backbone(x) |
| 31 | |
| 32 | if self.training: |
| 33 | assert targets is not None |
| 34 | loss, iou_loss, conf_loss, cls_loss, l1_loss, num_fg = self.head( |
| 35 | fpn_outs, targets, x |
| 36 | ) |
| 37 | outputs = { |
| 38 | "total_loss": loss, |
| 39 | "iou_loss": iou_loss, |
| 40 | "l1_loss": l1_loss, |
| 41 | "conf_loss": conf_loss, |
| 42 | "cls_loss": cls_loss, |
| 43 | "num_fg": num_fg, |
| 44 | } |
| 45 | else: |
| 46 | outputs = self.head(fpn_outs) |
| 47 | |
| 48 | return outputs |