Last commit not found
| # taken from https://raw.githubusercontent.com/janghyuncho/PiCIE/1d7b034f57e98670b0d6a244b2eea11fa0dde73e/modules/fpn.py | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from . import backbone_picie as backbone | |
| class PanopticFPN(nn.Module): | |
| def __init__(self, arch, pretrain, n_cls): | |
| super(PanopticFPN, self).__init__() | |
| self.n_cls = n_cls | |
| self.backbone = backbone.__dict__[arch](pretrained=pretrain) | |
| self.decoder = FPNDecoder(arch, n_cls) | |
| def forward(self, x, encoder_features=False, decoder_features=False): | |
| feats = self.backbone(x) | |
| if decoder_features: | |
| dec, outs = self.decoder(feats, get_features=decoder_features) | |
| else: | |
| outs = self.decoder(feats) | |
| if encoder_features: | |
| if decoder_features: | |
| return feats['res5'], dec, outs | |
| else: | |
| return feats['res5'], outs | |
| else: | |
| return outs | |
| class FPNDecoder(nn.Module): | |
| def __init__(self, arch, n_cls): | |
| super(FPNDecoder, self).__init__() | |
| self.n_cls = n_cls | |
| if arch == 'resnet18': | |
| mfactor = 1 | |
| out_dim = 128 | |
| else: | |
| mfactor = 4 | |
| out_dim = 256 | |
| self.layer4 = nn.Conv2d(512 * mfactor // 8, out_dim, kernel_size=1, stride=1, padding=0) | |
| self.layer3 = nn.Conv2d(512 * mfactor // 4, out_dim, kernel_size=1, stride=1, padding=0) | |
| self.layer2 = nn.Conv2d(512 * mfactor // 2, out_dim, kernel_size=1, stride=1, padding=0) | |
| self.layer1 = nn.Conv2d(512 * mfactor, out_dim, kernel_size=1, stride=1, padding=0) | |
| self.pred = nn.Conv2d(out_dim, self.n_cls, 1, 1) | |
| def forward(self, x, get_features=False): | |
| o1 = self.layer1(x['res5']) | |
| o2 = self.upsample_add(o1, self.layer2(x['res4'])) | |
| o3 = self.upsample_add(o2, self.layer3(x['res3'])) | |
| o4 = self.upsample_add(o3, self.layer4(x['res2'])) | |
| pred = self.pred(o4) | |
| if get_features: | |
| return o4, pred | |
| else: | |
| return pred | |
| def upsample_add(self, x, y): | |
| _, _, H, W = y.size() | |
| return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=False) + y | |