|
|
|
|
|
|
|
|
|
import os |
|
|
|
import torch.nn as nn |
|
|
|
from yolox.exp import Exp as MyExp |
|
|
|
|
|
class Exp(MyExp): |
|
def __init__(self): |
|
super(Exp, self).__init__() |
|
self.depth = 1.0 |
|
self.width = 1.0 |
|
self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] |
|
|
|
def get_model(self, sublinear=False): |
|
def init_yolo(M): |
|
for m in M.modules(): |
|
if isinstance(m, nn.BatchNorm2d): |
|
m.eps = 1e-3 |
|
m.momentum = 0.03 |
|
if "model" not in self.__dict__: |
|
from yolox.models import YOLOX, YOLOFPN, YOLOXHead |
|
backbone = YOLOFPN() |
|
head = YOLOXHead(self.num_classes, self.width, in_channels=[128, 256, 512], act="lrelu") |
|
self.model = YOLOX(backbone, head) |
|
self.model.apply(init_yolo) |
|
self.model.head.initialize_biases(1e-2) |
|
|
|
return self.model |
|
|
