From b3d7580df19d2ec7d6cb93b7a6a7f10c691abb25 Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Mon, 12 Apr 2021 12:11:35 -0700
Subject: [PATCH] Update ByoaNet comments. Fix first Steam feat chs for
 ByobNet.

---
 timm/models/byoanet.py | 20 +++++++++++---------
 timm/models/byobnet.py |  2 +-
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/timm/models/byoanet.py b/timm/models/byoanet.py
index faf366e6..7b7481dd 100644
--- a/timm/models/byoanet.py
+++ b/timm/models/byoanet.py
@@ -8,7 +8,7 @@ Currently used to implement experimential variants of:
   * Lambda ResNets
   * HaloNets
 
-Consider all of the models here a WIP and likely to change.
+Consider all of the models definitions here as experimental WIP and likely to change.
 
 Hacked together by / copyright Ross Wightman, 2021.
 """
@@ -372,7 +372,7 @@ def _create_byoanet(variant, cfg_variant=None, pretrained=False, **kwargs):
 
 @register_model
 def botnet50t_224(pretrained=False, **kwargs):
-    """
+    """ Bottleneck Transformer w/ ResNet50-T backbone. Bottleneck attn in final stage.
     """
     kwargs.setdefault('img_size', 224)
     return _create_byoanet('botnet50t_224', 'botnet50t', pretrained=pretrained, **kwargs)
@@ -380,7 +380,7 @@ def botnet50t_224(pretrained=False, **kwargs):
 
 @register_model
 def botnet50t_c4c5_224(pretrained=False, **kwargs):
-    """
+    """ Bottleneck Transformer w/ ResNet50-T backbone. Bottleneck attn in last two stages.
     """
     kwargs.setdefault('img_size', 224)
     return _create_byoanet('botnet50t_c4c5_224', 'botnet50t_c4c5', pretrained=pretrained, **kwargs)
@@ -388,41 +388,43 @@ def botnet50t_c4c5_224(pretrained=False, **kwargs):
 
 @register_model
 def halonet_h1(pretrained=False, **kwargs):
-    """
+    """ HaloNet-H1. Halo attention in all stages as per the paper.
+
+    This runs very slowly, param count lower than paper --> something is wrong.
     """
     return _create_byoanet('halonet_h1', pretrained=pretrained, **kwargs)
 
 
 @register_model
 def halonet_h1_c4c5(pretrained=False, **kwargs):
-    """
+    """ HaloNet-H1 config w/ attention in last two stages.
     """
     return _create_byoanet('halonet_h1_c4c5', pretrained=pretrained, **kwargs)
 
 
 @register_model
 def halonet26t(pretrained=False, **kwargs):
-    """
+    """ HaloNet w/ a ResNet26-t backbone, Hallo attention in final stage
     """
     return _create_byoanet('halonet26t', pretrained=pretrained, **kwargs)
 
 
 @register_model
 def halonet50t(pretrained=False, **kwargs):
-    """
+    """ HaloNet w/ a ResNet50-t backbone, Hallo attention in final stage
     """
     return _create_byoanet('halonet50t', pretrained=pretrained, **kwargs)
 
 
 @register_model
 def lambda_resnet26t(pretrained=False, **kwargs):
-    """
+    """ Lambda-ResNet-26T. Lambda layers in one C4 stage and all C5.
     """
     return _create_byoanet('lambda_resnet26t', pretrained=pretrained, **kwargs)
 
 
 @register_model
 def lambda_resnet50t(pretrained=False, **kwargs):
-    """
+    """ Lambda-ResNet-50T. Lambda layers in one C4 stage and all C5.
     """
     return _create_byoanet('lambda_resnet50t', pretrained=pretrained, **kwargs)
diff --git a/timm/models/byobnet.py b/timm/models/byobnet.py
index 2f6dbbef..75610f67 100644
--- a/timm/models/byobnet.py
+++ b/timm/models/byobnet.py
@@ -546,7 +546,7 @@ class Stem(nn.Sequential):
             layer_fn = layers.conv_norm_act if na else create_conv2d
             conv_name = f'conv{i + 1}'
             if i > 0 and s > 1:
-                self.feature_info.append(dict(num_chs=ch, reduction=curr_stride, module=prev_feat))
+                self.feature_info.append(dict(num_chs=prev_chs, reduction=curr_stride, module=prev_feat))
             self.add_module(conv_name, layer_fn(prev_chs, ch, kernel_size=kernel_size, stride=s))
             prev_chs = ch
             curr_stride *= s