chest-falsetto / model.py
monet-joe's picture
Upload 3 files
20e29fa verified
import torch
import torch.nn as nn
import torchvision.models as models
from modelscope.msdatasets import MsDataset
from utils import MODEL_DIR
def get_backbone(ver, backbone_list):
for bb in backbone_list:
if ver == bb["ver"]:
return bb
print("Backbone name not found, using default option - alexnet.")
return backbone_list[0]
def model_info(m_ver):
backbone_list = MsDataset.load(
"monetjoe/cv_backbones", split="train"
)
backbone = get_backbone(m_ver, backbone_list)
m_type = str(backbone["type"])
input_size = int(backbone["input_size"])
return m_type, input_size
def Classifier(cls_num: int, output_size: int, linear_output: bool):
q = (1.0 * output_size / cls_num) ** 0.25
l1 = int(q * cls_num)
l2 = int(q * l1)
l3 = int(q * l2)
if linear_output:
return torch.nn.Sequential(
nn.Dropout(),
nn.Linear(output_size, l3),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(l3, l2),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(l2, l1),
nn.ReLU(inplace=True),
nn.Linear(l1, cls_num),
)
else:
return torch.nn.Sequential(
nn.Dropout(),
nn.Conv2d(output_size, l3, kernel_size=(1, 1), stride=(1, 1)),
nn.ReLU(inplace=True),
nn.AdaptiveAvgPool2d(output_size=(1, 1)),
nn.Flatten(),
nn.Linear(l3, l2),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(l2, l1),
nn.ReLU(inplace=True),
nn.Linear(l1, cls_num),
)
class EvalNet:
model = None
m_type = "squeezenet"
input_size = 224
output_size = 512
def __init__(self, log_name: str, cls_num=4):
saved_model_path = f"{MODEL_DIR}/{log_name}/save.pt"
m_ver = "_".join(log_name.split("_")[:-1])
self.m_type, self.input_size = model_info(m_ver)
if not hasattr(models, m_ver):
print("Unsupported model.")
exit()
self.model = eval("models.%s()" % m_ver)
linear_output = self._set_outsize()
self._set_classifier(cls_num, linear_output)
checkpoint = torch.load(saved_model_path, map_location="cpu")
if torch.cuda.is_available():
checkpoint = torch.load(saved_model_path)
self.model.load_state_dict(checkpoint, False)
self.model.eval()
def _set_outsize(self, debug_mode=False):
for name, module in self.model.named_modules():
if (
str(name).__contains__("classifier")
or str(name).__eq__("fc")
or str(name).__contains__("head")
):
if isinstance(module, torch.nn.Linear):
self.output_size = module.in_features
if debug_mode:
print(
f"{name}(Linear): {self.output_size} -> {module.out_features}"
)
return True
if isinstance(module, torch.nn.Conv2d):
self.output_size = module.in_channels
if debug_mode:
print(
f"{name}(Conv2d): {self.output_size} -> {module.out_channels}"
)
return False
return False
def _set_classifier(self, cls_num, linear_output):
if self.m_type == "convnext":
del self.model.classifier[2]
self.model.classifier = nn.Sequential(
*list(self.model.classifier)
+ list(Classifier(cls_num, self.output_size, linear_output))
)
return
if hasattr(self.model, "classifier"):
self.model.classifier = Classifier(cls_num, self.output_size, linear_output)
return
elif hasattr(self.model, "fc"):
self.model.fc = Classifier(cls_num, self.output_size, linear_output)
return
elif hasattr(self.model, "head"):
self.model.head = Classifier(cls_num, self.output_size, linear_output)
return
self.model.heads.head = Classifier(cls_num, self.output_size, linear_output)
def forward(self, x):
if torch.cuda.is_available():
x = x.cuda()
self.model = self.model.cuda()
if self.m_type == "googlenet" and self.training:
return self.model(x)[0]
else:
return self.model(x)