Spaces:
Sleeping
Sleeping
| 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 | |