| 20 | self.enable_mixup = False |
| 21 | |
| 22 | def get_model(self, sublinear=False): |
| 23 | |
| 24 | def init_yolo(M): |
| 25 | for m in M.modules(): |
| 26 | if isinstance(m, nn.BatchNorm2d): |
| 27 | m.eps = 1e-3 |
| 28 | m.momentum = 0.03 |
| 29 | if "model" not in self.__dict__: |
| 30 | from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead |
| 31 | in_channels = [256, 512, 1024] |
| 32 | # NANO model use depthwise = True, which is main difference. |
| 33 | backbone = YOLOPAFPN(self.depth, self.width, in_channels=in_channels, depthwise=True) |
| 34 | head = YOLOXHead(self.num_classes, self.width, in_channels=in_channels, depthwise=True) |
| 35 | self.model = YOLOX(backbone, head) |
| 36 | |
| 37 | self.model.apply(init_yolo) |
| 38 | self.model.head.initialize_biases(1e-2) |
| 39 | return self.model |