| import torch |
| import torch.nn as nn |
| from transformers import PreTrainedModel, PretrainedConfig |
| from transformers.modeling_outputs import ImageClassifierOutput |
|
|
| class ResNetConfig(PretrainedConfig): |
| model_type = "custom_resnet" |
| def __init__(self, num_classes=10, num_channels=3, **kwargs): |
| super().__init__(**kwargs) |
| self.num_classes = num_classes |
| self.num_channels = num_channels |
|
|
| class BasicBlock(nn.Module): |
| expansion = 1 |
| def __init__(self, in_channels, out_channels, stride=1, downsample=None): |
| super(BasicBlock, self).__init__() |
| self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) |
| self.bn1 = nn.BatchNorm2d(out_channels) |
| self.relu = nn.ReLU(inplace=True) |
| self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) |
| self.bn2 = nn.BatchNorm2d(out_channels) |
| self.downsample = downsample |
|
|
| def forward(self, x): |
| identity = x |
| out = self.relu(self.bn1(self.conv1(x))) |
| out = self.bn2(self.conv2(out)) |
| if self.downsample is not None: |
| identity = self.downsample(x) |
| out += identity |
| return self.relu(out) |
|
|
| class ResNet(PreTrainedModel): |
| config_class = ResNetConfig |
| |
| def __init__(self, config): |
| super().__init__(config) |
| self.in_channels = 64 |
| self.conv1 = nn.Conv2d(config.num_channels, 64, kernel_size=7, stride=2, padding=3, bias=False) |
| self.bn1 = nn.BatchNorm2d(64) |
| self.relu = nn.ReLU(inplace=True) |
| self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) |
| |
| self.layer1 = self._make_layer(BasicBlock, 64, 3) |
| self.layer2 = self._make_layer(BasicBlock, 128, 4, stride=2) |
| self.layer3 = self._make_layer(BasicBlock, 256, 6, stride=2) |
| self.layer4 = self._make_layer(BasicBlock, 512, 3, stride=2) |
| |
| self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) |
| self.fc = nn.Linear(512, config.num_classes) |
|
|
| def _make_layer(self, block, out_channels, blocks, stride=1): |
| downsample = None |
| if stride != 1 or self.in_channels != out_channels: |
| downsample = nn.Sequential( |
| nn.Conv2d(self.in_channels, out_channels, kernel_size=1, stride=stride, bias=False), |
| nn.BatchNorm2d(out_channels), |
| ) |
| layers = [block(self.in_channels, out_channels, stride, downsample)] |
| self.in_channels = out_channels |
| for _ in range(1, blocks): |
| layers.append(block(self.in_channels, out_channels)) |
| return nn.Sequential(*layers) |
|
|
| def forward(self, pixel_values=None, labels=None, **kwargs): |
| x = pixel_values |
| x = self.relu(self.bn1(self.conv1(x))) |
| x = self.maxpool(x) |
| x = self.layer1(x) |
| x = self.layer2(x) |
| x = self.layer3(x) |
| x = self.layer4(x) |
| x = torch.flatten(self.avgpool(x), 1) |
| logits = self.fc(x) |
| |
| loss = None |
| if labels is not None: |
| loss_fct = nn.CrossEntropyLoss() |
| loss = loss_fct(logits, labels) |
| |
| return ImageClassifierOutput( |
| loss=loss, |
| logits=logits |
| ) |
