added model.py

model.py

@@ -0,0 +1,81 @@
+import torch
+
+def model_init(m):
+    if isinstance(m, torch.nn.Conv2d) or isinstance(m, torch.nn.Linear):
+        torch.nn.init.xavier_uniform_(m.weight)
+        torch.nn.init.zeros_(m.bias)
+
+class NLB(torch.nn.Module):
+    def __init__(self, in_ch, relu_a=0.01):
+        self.inter_ch = torch.div(in_ch, 2, rounding_mode='floor').item()
+        super().__init__()
+        self.theta_layer = torch.nn.Conv2d(in_channels=in_ch, out_channels=self.inter_ch, \
+                            kernel_size=1, padding=0)
+        self.phi_layer   = torch.nn.Conv2d(in_channels=in_ch, out_channels=self.inter_ch, \
+                            kernel_size=1, padding=0)
+        self.g_layer     = torch.nn.Conv2d(in_channels=in_ch, out_channels=self.inter_ch, \
+                            kernel_size=1, padding=0)
+        self.atten_act   = torch.nn.Softmax(dim=-1)
+        self.out_cnn     = torch.nn.Conv2d(in_channels=self.inter_ch, out_channels=in_ch, \
+                            kernel_size=1, padding=0)
+        
+    def forward(self, x):
+        mbsz, _, h, w = x.size()
+        
+        theta = self.theta_layer(x).view(mbsz, self.inter_ch, -1).permute(0, 2, 1)
+        phi   = self.phi_layer(x).view(mbsz, self.inter_ch, -1)
+        g     = self.g_layer(x).view(mbsz, self.inter_ch, -1).permute(0, 2, 1)
+        
+        theta_phi = self.atten_act(torch.matmul(theta, phi))
+        
+        theta_phi_g = torch.matmul(theta_phi, g).permute(0, 2, 1).view(mbsz, self.inter_ch, h, w)
+        
+        _out_tmp = self.out_cnn(theta_phi_g)
+        _out_tmp = torch.add(_out_tmp, x)
+   
+        return _out_tmp
+
+
+class BraggNN(torch.nn.Module):
+    def __init__(self, imgsz, fcsz=(64, 32, 16, 8)):
+        super().__init__()
+        self.cnn_ops = []
+        cnn_out_chs = (64, 32, 8)
+        cnn_in_chs  = (1, ) + cnn_out_chs[:-1]
+        fsz = imgsz
+        for ic, oc, in zip(cnn_in_chs, cnn_out_chs):
+            self.cnn_ops += [
+                            torch.nn.Conv2d(in_channels=ic, out_channels=oc, kernel_size=3, \
+                                stride=1, padding=0),
+                            torch.nn.LeakyReLU(negative_slope=0.01),
+            ]
+            fsz -= 2
+        self.nlb = NLB(in_ch=cnn_out_chs[0])
+        self.dense_ops = []
+        dense_in_chs  = (fsz * fsz * cnn_out_chs[-1], ) + fcsz[:-1]
+        for ic, oc in zip(dense_in_chs, fcsz):
+            self.dense_ops += [
+                            torch.nn.Linear(ic, oc),
+                            torch.nn.LeakyReLU(negative_slope=0.01),
+            ]
+        # output layer
+        self.dense_ops += [torch.nn.Linear(fcsz[-1], 2), ]
+        
+        self.cnn_layers   = torch.nn.Sequential(*self.cnn_ops)
+        self.dense_layers = torch.nn.Sequential(*self.dense_ops)
+        
+    def forward(self, x):
+        _out = x
+        for layer in self.cnn_layers[:1]:
+            _out = layer(_out)
+
+        _out = self.nlb(_out)
+
+        for layer in self.cnn_layers[1:]:
+            _out = layer(_out)
+        
+        _out = _out.flatten(start_dim=1)
+        for layer in self.dense_layers:
+            _out = layer(_out)
+            
+        return _out