File size: 2,068 Bytes
c985ba4 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from networks.layers.basic import ConvGN
class FPNSegmentationHead(nn.Module):
def __init__(self,
in_dim,
out_dim,
decode_intermediate_input=True,
hidden_dim=256,
shortcut_dims=[24, 32, 96, 1280],
align_corners=True):
super().__init__()
self.align_corners = align_corners
self.decode_intermediate_input = decode_intermediate_input
self.conv_in = ConvGN(in_dim, hidden_dim, 1)
self.conv_16x = ConvGN(hidden_dim, hidden_dim, 3)
self.conv_8x = ConvGN(hidden_dim, hidden_dim // 2, 3)
self.conv_4x = ConvGN(hidden_dim // 2, hidden_dim // 2, 3)
self.adapter_16x = nn.Conv2d(shortcut_dims[-2], hidden_dim, 1)
self.adapter_8x = nn.Conv2d(shortcut_dims[-3], hidden_dim, 1)
self.adapter_4x = nn.Conv2d(shortcut_dims[-4], hidden_dim // 2, 1)
self.conv_out = nn.Conv2d(hidden_dim // 2, out_dim, 1)
self._init_weight()
def forward(self, inputs, shortcuts):
if self.decode_intermediate_input:
x = torch.cat(inputs, dim=1)
else:
x = inputs[-1]
x = F.relu_(self.conv_in(x))
x = F.relu_(self.conv_16x(self.adapter_16x(shortcuts[-2]) + x))
x = F.interpolate(x,
size=shortcuts[-3].size()[-2:],
mode="bilinear",
align_corners=self.align_corners)
x = F.relu_(self.conv_8x(self.adapter_8x(shortcuts[-3]) + x))
x = F.interpolate(x,
size=shortcuts[-4].size()[-2:],
mode="bilinear",
align_corners=self.align_corners)
x = F.relu_(self.conv_4x(self.adapter_4x(shortcuts[-4]) + x))
x = self.conv_out(x)
return x
def _init_weight(self):
for p in self.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
|