Update qlnet.py

qlnet.py

@@ -19,6 +19,8 @@ from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 from timm.layers import DropBlock2d, DropPath, AvgPool2dSame, BlurPool2d, GroupNorm, create_attn, get_attn, \
     get_act_layer, get_norm_layer, create_classifier, LayerNorm2d
 
+from ._builder import build_model_with_cfg
+from ._registry import register_model, model_entrypoint
 
 def get_padding(kernel_size, stride, dilation=1):
     padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
@@ -104,8 +106,8 @@ class QLBlock(nn.Module): # quasilinear hyperbolic system
     ):
         super(QLBlock, self).__init__()
 
-        k = 4 if inplanes <= 256 else 2
-        width = inplanes * k
+        self.k = 8 if inplanes <= 128 else 4 if inplanes <= 256 else 2
+        width = inplanes * self.k
         outplanes = inplanes if downsample is None else inplanes * 2
         first_dilation = first_dilation or dilation
 
@@ -114,6 +116,8 @@ class QLBlock(nn.Module): # quasilinear hyperbolic system
                 dilation=first_dilation, groups=1, bias=False),
             norm_layer(width*2))
 
+        # self.conv2 = nn.Conv2d(1, self.k, kernel_size=3, stride=stride,
+        #         padding=1, dilation=first_dilation, groups=1, bias=False)
         self.conv2 = nn.Conv2d(width, width, kernel_size=3, stride=stride,
                 padding=1, dilation=first_dilation, groups=width, bias=False)
         self.bn2 = norm_layer(width)
@@ -135,18 +139,17 @@ class QLBlock(nn.Module): # quasilinear hyperbolic system
 
     def conv_forward(self, x):
         conv = self.conv2
-        k = conv.in_channels
-        C = x.size()[1] // k
+        C = x.size(1) // self.k
         kernel = conv.weight.repeat(C, 1, 1, 1)
         bias = conv.bias.repeat(C) if conv.bias is not None else None
         return F.conv2d(x, kernel, bias, conv.stride, 
-            conv.padding, conv.dilation, C * k)
+            conv.padding, conv.dilation, x.size(1))
 
     def forward(self, x):
         x0 = self.skip(x)
         x = self.conv1(x)
-        C = x.size(1) // 2
-        x = x[:, :C, :, :] * x[:, C:, :, :]
+        x = x[:, ::2, :, :] * x[:, 1::2, :, :]
+        
         x = self.conv2(x)
         x = self.bn2(x)
         x = self.conv3(x)
@@ -209,15 +212,15 @@ class QLNet(nn.Module):
     def __init__(
             self,
             block=QLBlock,   # new block
-            layers=[3,4,6,3], # as in resnet50
+            layers=[3,4,12,3], # [3,4,6,3] as in resnet50
             num_classes=1000,
             in_chans=3,
             output_stride=32,
             global_pool='avg',
             cardinality=1,
             base_width=64,
-            stem_width=64,
-            stem_type='',
+            stem_width=32,
+            stem_type='', # 'deep' for resnet50d
             replace_stem_pool=False,
             block_reduce_first=1,
             down_kernel_size=1,
@@ -280,7 +283,7 @@ class QLNet(nn.Module):
         else:
             self.conv1 = nn.Conv2d(in_chans, inplanes, kernel_size=7, stride=2, padding=3, bias=False)
         self.bn1 = norm_layer(inplanes)
-        self.act1 = act_layer(inplace=True)
+        # self.act1 = act_layer(inplace=True)
         self.feature_info = [dict(num_chs=inplanes, reduction=2, module='act1')]
 
         # Stem pooling. The name 'maxpool' remains for weight compatibility.
@@ -314,7 +317,7 @@ class QLNet(nn.Module):
             self.add_module(*stage)  # layer1, layer2, etc
         self.feature_info.extend(stage_feature_info)
 
-        self.act = hardball(radius2=512)
+        # self.act = hardball(radius2=512)
         # self.act = nn.Hardtanh(max_val=5, min_val=-5, inplace=True)
         # self.act = nn.ReLU(inplace=True)
 
@@ -361,7 +364,7 @@ class QLNet(nn.Module):
     def forward_features(self, x):
         x = self.conv1(x)
         x = self.bn1(x)
-        x = self.act1(x)
+        # x = self.act1(x)
         x = self.maxpool(x)
 
         if self.grad_checkpointing and not torch.jit.is_scripting():
@@ -379,7 +382,26 @@ class QLNet(nn.Module):
 
     def forward(self, x):
         x = self.forward_features(x)
-        x = self.act(x)
+        # x = self.act(x)
         x = self.forward_head(x)
         return x
 
+
+def _create_qlnet(variant, pretrained=False, **kwargs):
+    return build_model_with_cfg(QLNet, variant, pretrained, **kwargs)
+
+
+@register_model
+def qlnet22(pretrained=False, **kwargs):
+    """Constructs a QLNet22 model.
+    """
+    model_args = dict(block=QLBlock, layers=[3, 4, 12, 3],  **kwargs)
+    return _create_qlnet('qlnet22', pretrained, **dict(model_args, **kwargs))
+
+
+@register_model
+def qlnet26(pretrained=False, **kwargs):
+    """Constructs a QLNet26 model.
+    """
+    model_args = dict(block=QLBlock, layers=[3, 8, 12, 3],  **kwargs)
+    return _create_qlnet('qlnet26', pretrained, **dict(model_args, **kwargs))