From 5ba9390a7e4e8dbf366cfd280403cf37a6b22bab Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Wed, 30 Dec 2020 16:58:17 +0800 Subject: Combine FPFE and TFA to GaitPart --- models/layers.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) (limited to 'models/layers.py') diff --git a/models/layers.py b/models/layers.py index 62a3cc6..c69ae07 100644 --- a/models/layers.py +++ b/models/layers.py @@ -119,6 +119,33 @@ class FocalConv2d(BasicConv2d): return F.leaky_relu(z, inplace=True) +class FocalConv2dBlock(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_sizes: tuple[int, int], + paddings: tuple[int, int], + halving: int, + use_pool: bool = True, + **kwargs + ): + super().__init__() + self.use_pool = use_pool + self.fconv1 = FocalConv2d(in_channels, out_channels, kernel_sizes[0], + halving, padding=paddings[0], **kwargs) + self.fconv2 = FocalConv2d(out_channels, out_channels, kernel_sizes[1], + halving, padding=paddings[1], **kwargs) + self.max_pool = nn.MaxPool2d(2) + + def forward(self, x): + x = self.fconv1(x) + x = self.fconv2(x) + if self.use_pool: + x = self.max_pool(x) + return x + + class BasicConv1d(nn.Module): def __init__( self, -- cgit v1.2.3