Spaces:
Sleeping
Sleeping
File size: 2,159 Bytes
749745d |
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 |
from torch import nn
from maskrcnn_benchmark.modeling.make_layers import make_conv3x3
class Residual(nn.Module):
def __init__(self, inp_dim, out_dim, use_gn=False):
super(Residual, self).__init__()
self.relu = nn.ReLU()
# self.bn1 = nn.BatchNorm2d(inp_dim)
self.conv1 = make_conv3x3(inp_dim, int(out_dim / 2), 1, use_relu=False, use_gn=use_gn)
# self.bn2 = nn.BatchNorm2d(int(out_dim / 2))
self.conv2 = make_conv3x3(int(out_dim / 2), int(out_dim / 2), 3, use_relu=False, use_gn=use_gn)
# self.bn3 = nn.BatchNorm2d(int(out_dim / 2))
self.conv3 = make_conv3x3(int(out_dim / 2), out_dim, 1, use_relu=False, use_gn=use_gn)
if inp_dim == out_dim:
self.need_skip = False
else:
self.need_skip = True
self.skip_layer = make_conv3x3(inp_dim, out_dim, 1, use_relu=False, use_gn=False)
def forward(self, x):
if self.need_skip:
residual = self.skip_layer(x)
else:
residual = x
out = x
# out = self.bn1(out)
out = self.relu(out)
out = self.conv1(out)
# out = self.bn2(out)
out = self.relu(out)
out = self.conv2(out)
# out = self.bn3(out)
out = self.relu(out)
out = self.conv3(out)
out += residual
return out
class Hourglass(nn.Module):
def __init__(self, n, f, gn=False, increase=0):
super(Hourglass, self).__init__()
nf = f + increase
self.up1 = Residual(f, f)
# Lower branch
self.pool1 = nn.MaxPool2d(2, 2)
self.low1 = Residual(f, nf)
self.n = n
# Recursive hourglass
if self.n > 1:
self.low2 = Hourglass(n - 1, nf, gn=gn)
else:
self.low2 = Residual(nf, nf, gn)
self.low3 = Residual(nf, f, gn)
self.up2 = nn.Upsample(scale_factor=2, mode="nearest")
def forward(self, x):
up1 = self.up1(x)
pool1 = self.pool1(x)
low1 = self.low1(pool1)
low2 = self.low2(low1)
low3 = self.low3(low2)
up2 = self.up2(low3)
return up1 + up2
|