From 96f345d25237c7e616ea5f524a2fc2d340ed8aff Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Wed, 23 Dec 2020 18:59:08 +0800
Subject: Split modules to different files

---
 models/layers.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)
 create mode 100644 models/layers.py

(limited to 'models/layers.py')

diff --git a/models/layers.py b/models/layers.py
new file mode 100644
index 0000000..e737df2
--- /dev/null
+++ b/models/layers.py
@@ -0,0 +1,27 @@
+from typing import Union, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class FocalConv2d(nn.Module):
+    def __init__(
+            self,
+            in_channels: int,
+            out_channels: int,
+            kernel_size: Union[int, Tuple[int, int]],
+            halving: int,
+            **kwargs
+    ):
+        super(FocalConv2d, self).__init__()
+        self.halving = halving
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,
+                              bias=False, **kwargs)
+
+    def forward(self, x):
+        h = x.size(2)
+        split_size = h // 2 ** self.halving
+        z = x.split(split_size, dim=2)
+        z = torch.cat([self.conv(_) for _ in z], dim=2)
+        return F.leaky_relu(z, inplace=True)
-- 
cgit v1.2.3