File size: 4,919 Bytes
309069b 844c88e 309069b 7f6490f 309069b 7f6490f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import ImageClassifierOutput
from .configuration_mobilenet import MobileNetV1Config
class FP32LayerNorm2d(nn.GroupNorm):
def __init__(self, num_channels):
super().__init__(1, num_channels)
def forward(self, x):
input_dtype = x.dtype
with torch.autocast(device_type=x.device.type, enabled=False):
normalized = super().forward(x.float())
return normalized.to(dtype=input_dtype)
class DepthwiseSeparableConv(nn.Module):
def __init__(self, in_channels, out_channels, stride, dropout=0.0):
super().__init__()
self.use_residual = (stride == 1 and in_channels == out_channels)
self.dw = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=stride, padding=1, groups=in_channels, bias=False)
self.dw_norm = FP32LayerNorm2d(in_channels)
self.dw_act = nn.SiLU(inplace=True)
self.pw = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False)
self.pw_norm = FP32LayerNorm2d(out_channels)
self.pw_act = nn.SiLU(inplace=True)
self.dropout = nn.Dropout2d(p=dropout) if dropout > 0 else nn.Identity()
if self.use_residual:
self.residual_scale = nn.Parameter(torch.tensor(0.1))
def forward(self, x):
identity = x
out = self.dw(x)
out = self.dw_norm(out)
out = self.dw_act(out)
out = self.pw(out)
out = self.pw_norm(out)
out = self.pw_act(out)
out = self.dropout(out)
if self.use_residual:
return identity + self.residual_scale * out
return out
class MobileNetV1ForImageClassification(PreTrainedModel):
config_class = MobileNetV1Config
def __init__(self, config: MobileNetV1Config):
super().__init__(config)
self.num_labels = config.num_classes
self.stem = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False),
FP32LayerNorm2d(32),
nn.SiLU(inplace=True)
)
self.blocks = nn.Sequential(
DepthwiseSeparableConv(32, 64, stride=1, dropout=config.block_dropout),
DepthwiseSeparableConv(64, 128, stride=2, dropout=config.block_dropout),
DepthwiseSeparableConv(128, 128, stride=1, dropout=config.block_dropout),
DepthwiseSeparableConv(128, 256, stride=2, dropout=config.block_dropout),
DepthwiseSeparableConv(256, 256, stride=1, dropout=config.block_dropout),
DepthwiseSeparableConv(256, 512, stride=2, dropout=config.block_dropout),
DepthwiseSeparableConv(512, 512, stride=1, dropout=config.block_dropout),
DepthwiseSeparableConv(512, 512, stride=1, dropout=config.block_dropout),
DepthwiseSeparableConv(512, 512, stride=1, dropout=config.block_dropout),
DepthwiseSeparableConv(512, 512, stride=1, dropout=config.block_dropout),
DepthwiseSeparableConv(512, 512, stride=1, dropout=config.block_dropout),
DepthwiseSeparableConv(512, 1024, stride=2, dropout=config.block_dropout),
DepthwiseSeparableConv(1024, 1024, stride=1, dropout=config.block_dropout),
)
self.gap = nn.AdaptiveAvgPool2d((1, 1))
self.dropout = nn.Dropout(p=config.final_dropout)
self.classifier = nn.Linear(1024, config.num_classes)
self.post_init()
def _init_weights(self, module):
if isinstance(module, nn.Conv2d):
nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.GroupNorm):
nn.init.ones_(module.weight)
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Linear):
nn.init.normal_(module.weight, mean=0.0, std=0.001)
if module.bias is not None:
nn.init.zeros_(module.bias)
def forward(self, pixel_values=None, labels=None, return_dict=None):
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
x = self.stem(pixel_values)
x = self.blocks(x)
x = self.gap(x)
x = torch.flatten(x, 1)
x = self.dropout(x)
logits = self.classifier(x)
loss = None
if labels is not None:
loss_fct = nn.CrossEntropyLoss()
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
if not return_dict:
output = (logits,)
return ((loss,) + output) if loss is not None else output
return ImageClassifierOutput(
loss=loss,
logits=logits,
) |