Spaces:
Running on Zero
Running on Zero
File size: 2,034 Bytes
7405e85 | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
# Ref: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py#L363
class ViTFeaturePyramid(nn.Module):
"""
This module implements SimpleFeaturePyramid in :paper:`vitdet`.
It creates pyramid features built on top of the input feature map.
"""
def __init__(
self,
in_channels,
scale_factors,
):
"""
Args:
scale_factors (list[float]): list of scaling factors to upsample or downsample
the input features for creating pyramid features.
"""
super(ViTFeaturePyramid, self).__init__()
self.scale_factors = scale_factors
out_dim = dim = in_channels
self.stages = nn.ModuleList()
for idx, scale in enumerate(scale_factors):
if scale == 4.0:
layers = [
nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2),
nn.GELU(),
nn.ConvTranspose2d(dim // 2, dim // 4, kernel_size=2, stride=2),
]
out_dim = dim // 4
elif scale == 2.0:
layers = [nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2)]
out_dim = dim // 2
elif scale == 1.0:
layers = []
elif scale == 0.5:
layers = [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
raise NotImplementedError(f"scale_factor={scale} is not supported yet.")
if scale != 1.0:
layers.extend(
[
nn.GELU(),
nn.Conv2d(out_dim, out_dim, 3, 1, 1),
]
)
layers = nn.Sequential(*layers)
self.stages.append(layers)
def forward(self, x):
results = []
for stage in self.stages:
results.append(stage(x))
return results
|