Spaces:
Running
Running
File size: 1,293 Bytes
e02ffe6 |
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 |
import torch.nn as nn
import torch.nn.functional as F
from .repvgg import create_RepVGG
class RepVGG_8_1_align(nn.Module):
"""
RepVGG backbone, output resolution are 1/8 and 1.
Each block has 2 layers.
"""
def __init__(self, config):
super().__init__()
backbone = create_RepVGG(False)
self.layer0, self.layer1, self.layer2, self.layer3 = backbone.stage0, backbone.stage1, backbone.stage2, backbone.stage3
for layer in [self.layer0, self.layer1, self.layer2, self.layer3]:
for m in layer.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
def forward(self, x):
out = self.layer0(x) # 1/2
for module in self.layer1:
out = module(out) # 1/2
x1 = out
for module in self.layer2:
out = module(out) # 1/4
x2 = out
for module in self.layer3:
out = module(out) # 1/8
x3 = out
return {'feats_c': x3, 'feats_f': None, 'feats_x2': x2, 'feats_x1': x1}
|