import torch
import torch.nn as nn
from core.registry import Registry
from core.trainer import _initialize_weights
from .ml_heads import build_ml_head
from .layers import LocalBlock
MODEL_REGISTRY = Registry("MODEL_TYPE")
MODEL_REGISTRY.__doc__ = """
Registry for Mammo models.
"""
BACKBONE_REGISTRY = Registry("BACKBONE_TYPE")
BACKBONE_REGISTRY.__doc__ = """
Registry for Mammo cls backbone.
"""
def build_model(cfg):
"""Build the whole model architecture
"""
model_name = cfg.MODEL.NAME
model = MODEL_REGISTRY.get(model_name)(cfg)
return model
@MODEL_REGISTRY.register()
class ResNet_v0(nn.Module):
"""Original ResNet
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from core.registry import Registry
ML_HEAD_REGISTRY = Registry("ML_HEAD")
ML_HEAD_REGISTRY.__doc__ = """
Registry for metric learning heads.
"""
def build_ml_head(cfg, in_channels):
head_name = cfg.MODEL.ML_HEAD.NAME
out_channels = cfg.MODEL.NUM_CLASSES
s = cfg.MODEL.ML_HEAD.SCALER
m = cfg.MODEL.ML_HEAD.MARGIN
num_centers = cfg.MODEL.ML_HEAD.NUM_CENTERS
head = ML_HEAD_REGISTRY.get(head_name)(in_channels, out_channels, s, m,
num_centers)
return head
@ML_HEAD_REGISTRY.register()
class ArcFace(nn.Module):
"""
This module implements ArcFace.
"""
def __init__(self, in_channels, out_channels, s=32., m=0.5, num_centers=1):
super(ArcFace, self).__init__()