Spaces:
Running on Zero
Running on Zero
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import torch | |
| from torch import nn | |
| import torch.nn.functional as F | |
| class ResidualBlock(nn.Module): | |
| def __init__(self, in_planes: int, planes: int, stride: int = 1) -> None: | |
| super().__init__() | |
| self.conv1 = nn.Conv2d( | |
| in_planes, | |
| planes, | |
| kernel_size=3, | |
| padding=1, | |
| stride=stride, | |
| padding_mode="zeros", | |
| ) | |
| self.conv2 = nn.Conv2d( | |
| planes, | |
| planes, | |
| kernel_size=3, | |
| padding=1, | |
| padding_mode="zeros", | |
| ) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.norm1 = nn.InstanceNorm2d(planes) | |
| self.norm2 = nn.InstanceNorm2d(planes) | |
| if stride != 1: | |
| self.norm3 = nn.InstanceNorm2d(planes) | |
| if stride == 1: | |
| self.downsample = None | |
| else: | |
| self.downsample = nn.Sequential( | |
| nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), | |
| self.norm3, | |
| ) | |
| def forward(self, inputs: torch.Tensor) -> torch.Tensor: | |
| residual = inputs | |
| outputs = self.relu(self.norm1(self.conv1(inputs))) | |
| outputs = self.relu(self.norm2(self.conv2(outputs))) | |
| if self.downsample is not None: | |
| residual = self.downsample(residual) | |
| return self.relu(residual + outputs) | |
| class BasicEncoder(nn.Module): | |
| def __init__(self, input_dim=3, output_dim=128, stride=4): | |
| super().__init__() | |
| self.stride = stride | |
| self.in_planes = output_dim // 2 | |
| self.norm1 = nn.InstanceNorm2d(self.in_planes) | |
| self.norm2 = nn.InstanceNorm2d(output_dim * 2) | |
| self.conv1 = nn.Conv2d( | |
| input_dim, | |
| self.in_planes, | |
| kernel_size=7, | |
| stride=2, | |
| padding=3, | |
| padding_mode="zeros", | |
| ) | |
| self.relu1 = nn.ReLU(inplace=True) | |
| self.layer1 = self._make_layer(output_dim // 2, stride=1) | |
| self.layer2 = self._make_layer(output_dim // 4 * 3, stride=2) | |
| self.layer3 = self._make_layer(output_dim, stride=2) | |
| self.layer4 = self._make_layer(output_dim, stride=2) | |
| self.conv2 = nn.Conv2d( | |
| output_dim * 3 + output_dim // 4, | |
| output_dim * 2, | |
| kernel_size=3, | |
| padding=1, | |
| padding_mode="zeros", | |
| ) | |
| self.relu2 = nn.ReLU(inplace=True) | |
| self.conv3 = nn.Conv2d(output_dim * 2, output_dim, kernel_size=1) | |
| for module in self.modules(): | |
| if isinstance(module, nn.Conv2d): | |
| nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu") | |
| elif isinstance(module, nn.InstanceNorm2d): | |
| if module.weight is not None: | |
| nn.init.constant_(module.weight, 1) | |
| if module.bias is not None: | |
| nn.init.constant_(module.bias, 0) | |
| def _make_layer(self, dim, stride=1): | |
| layer1 = ResidualBlock(self.in_planes, dim, stride=stride) | |
| layer2 = ResidualBlock(dim, dim, stride=1) | |
| self.in_planes = dim | |
| return nn.Sequential(layer1, layer2) | |
| def forward(self, inputs: torch.Tensor) -> torch.Tensor: | |
| height, width = inputs.shape[-2:] | |
| features = self.relu1(self.norm1(self.conv1(inputs))) | |
| level1 = self.layer1(features) | |
| level2 = self.layer2(level1) | |
| level3 = self.layer3(level2) | |
| level4 = self.layer4(level3) | |
| output_shape = (height // self.stride, width // self.stride) | |
| levels = [ | |
| F.interpolate( | |
| level, | |
| output_shape, | |
| mode="bilinear", | |
| align_corners=True, | |
| ) | |
| for level in (level1, level2, level3, level4) | |
| ] | |
| features = self.relu2(self.norm2(self.conv2(torch.cat(levels, dim=1)))) | |
| return self.conv3(features) | |