File size: 896 Bytes
8b79d57 |
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 |
from torch import nn
class LRASPP(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.aspp1 = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(True)
)
self.aspp2 = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Conv2d(in_channels, out_channels, 1, bias=False),
nn.Sigmoid()
)
def forward_single_frame(self, x):
return self.aspp1(x) * self.aspp2(x)
def forward_time_series(self, x):
B, T = x.shape[:2]
x = self.forward_single_frame(x.flatten(0, 1)).unflatten(0, (B, T))
return x
def forward(self, x):
if x.ndim == 5:
return self.forward_time_series(x)
else:
return self.forward_single_frame(x) |