| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import PreTrainedModel |
|
|
| from .configuration_vehicle_encoder import VehicleEncoderConfig |
|
|
|
|
| class ResBlock(nn.Module): |
| def __init__(self, in_ch, out_ch, stride=1): |
| super().__init__() |
| self.conv1 = nn.Conv2d(in_ch, out_ch, 3, stride, 1, bias=False) |
| self.bn1 = nn.BatchNorm2d(out_ch) |
| self.conv2 = nn.Conv2d(out_ch, out_ch, 3, 1, 1, bias=False) |
| self.bn2 = nn.BatchNorm2d(out_ch) |
| if stride != 1 or in_ch != out_ch: |
| self.skip = nn.Sequential( |
| nn.Conv2d(in_ch, out_ch, 1, stride, bias=False), |
| nn.BatchNorm2d(out_ch), |
| ) |
| else: |
| self.skip = nn.Identity() |
|
|
| def forward(self, x): |
| identity = self.skip(x) |
| out = F.relu(self.bn1(self.conv1(x)), inplace=True) |
| out = self.bn2(self.conv2(out)) |
| return F.relu(out + identity, inplace=True) |
|
|
|
|
| def _make_stage(in_ch, out_ch, blocks, stride): |
| layers = [ResBlock(in_ch, out_ch, stride=stride)] |
| for _ in range(blocks - 1): |
| layers.append(ResBlock(out_ch, out_ch, stride=1)) |
| return nn.Sequential(*layers) |
|
|
|
|
| class VehicleEncoderModel(PreTrainedModel): |
| config_class = VehicleEncoderConfig |
|
|
| def __init__(self, config: VehicleEncoderConfig): |
| super().__init__(config) |
| bps = config.blocks_per_stage |
|
|
| self.stem = nn.Sequential( |
| nn.Conv2d(3, 64, 7, 2, 3, bias=False), |
| nn.BatchNorm2d(64), |
| nn.ReLU(inplace=True), |
| nn.MaxPool2d(3, 2, 1), |
| ) |
| self.stage1 = _make_stage( 64, 64, bps, stride=1) |
| self.stage2 = _make_stage( 64, 128, bps, stride=2) |
| self.stage3 = _make_stage(128, 256, bps, stride=2) |
| self.stage4 = _make_stage(256, 512, bps, stride=2) |
| self.gap = nn.AdaptiveAvgPool2d(1) |
| self.fc = nn.Linear(512, config.latent_dim) |
|
|
| self.post_init() |
|
|
| def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: |
| x = self.stem(pixel_values) |
| x = self.stage1(x); x = self.stage2(x) |
| x = self.stage3(x); x = self.stage4(x) |
| x = self.gap(x).flatten(1) |
| return self.fc(x) |
