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)
|
|
