update lsnet_artist.py in preparation of lsnet_xl_artist_448 arch

lsnet/lsnet_artist.py

@@ -1,11 +1,25 @@
+"""
+LSNet for Artist Style Classification and Clustering
+支持画师风格的分类和聚类任务
+"""
 import torch
 import torch.nn as nn
-from .lsnet import LSNet, Conv2d_BN, BN_Linear
-from timm.models import register_model
-from timm.models import build_model_with_cfg
+from timm.models import build_model_with_cfg, register_model
+
+from .lsnet import BN_Linear, Conv2d_BN, LSNet
 
 
 class LSNetArtist(LSNet):
+    """
+    LSNet模型用于画师风格分类和聚类
+    
+    特点:
+    - 训练时使用分类头进行监督学习
+    - 推理时可选择是否使用分类头
+    - 去掉分类头输出特征向量用于聚类
+    - 保留分类头可以做风格分类
+    """
+    
     def __init__(self, 
                  img_size=224,
                  patch_size=8,
@@ -71,14 +85,20 @@ class LSNetArtist(LSNet):
         x = self.projection(x)
         return x
     
-    def forward(self, x, return_features=False):
+    def forward(self, x, return_features=False, return_both=False):
         """
-        x: 输入图像
-        return_features: 是否只返回特征向量（用于聚类）
-                        False时返回分类logits（用于分类）
+        前向传播
         
-        如果return_features=True: 返回特征向量 (batch_size, feature_dim)
-        如果return_features=False: 返回分类logits (batch_size, num_classes)
+        Args:
+            x: 输入图像
+            return_features: 是否只返回特征向量（用于聚类）
+                           False时返回分类logits（用于分类）
+            return_both: 是否同时返回特征向量和分类logits（用于对比损失）
+        
+        Returns:
+            如果return_features=True: 返回特征向量 (batch_size, feature_dim)
+            如果return_both=True: 返回 (features, logits)
+            如果return_features=False and return_both=False: 返回分类logits (batch_size, num_classes)
         """
         features = self.forward_features(x)
         
@@ -88,23 +108,26 @@ class LSNetArtist(LSNet):
         
         # 返回分类结果
         if self.distillation:
-            x = self.head(features), self.head_dist(features)
+            logits = self.head(features), self.head_dist(features)
             if not self.training:
-                x = (x[0] + x[1]) / 2
+                logits = (logits[0] + logits[1]) / 2
         else:
-            x = self.head(features)
+            logits = self.head(features)
         
-        return x
+        if return_both:
+            return features, logits
+        
+        return logits
     
     def get_features(self, x):
         """
-        提取特征向量
+        便捷方法：提取特征向量
         """
         return self.forward(x, return_features=True)
     
     def classify(self, x):
         """
-        进行分类
+        便捷方法：进行分类
         """
         return self.forward(x, return_features=False)
 
@@ -129,8 +152,9 @@ default_cfgs_artist = dict(
     lsnet_t_artist = _cfg_artist(),
     lsnet_s_artist = _cfg_artist(),
     lsnet_b_artist = _cfg_artist(),
-    lsnet_l_artist = _cfg_artist(),
-    lsnet_xl_artist = _cfg_artist(),
+    lsnet_l_artist = _cfg_artist(),  # Large model for massive training
+    lsnet_xl_artist = _cfg_artist(), # Extra Large model for 100k+ classes
+    lsnet_xl_artist_448 = _cfg_artist(), # Extra Large model with 448x448 input for massive datasets with 50k+ classes
 )
 
 
@@ -151,6 +175,7 @@ def _create_lsnet_artist(variant, pretrained=False, **kwargs):
 @register_model
 def lsnet_t_artist(num_classes=1000, distillation=False, pretrained=False, 
                    feature_dim=None, use_projection=True, **kwargs):
+    """LSNet-T for Artist Style Classification"""
     model = _create_lsnet_artist(
         "lsnet_t_artist",
         pretrained=pretrained,
@@ -171,6 +196,7 @@ def lsnet_t_artist(num_classes=1000, distillation=False, pretrained=False,
 @register_model
 def lsnet_s_artist(num_classes=1000, distillation=False, pretrained=False,
                    feature_dim=None, use_projection=True, **kwargs):
+    """LSNet-S for Artist Style Classification"""
     model = _create_lsnet_artist(
         "lsnet_s_artist",
         pretrained=pretrained,
@@ -191,6 +217,7 @@ def lsnet_s_artist(num_classes=1000, distillation=False, pretrained=False,
 @register_model
 def lsnet_b_artist(num_classes=1000, distillation=False, pretrained=False,
                    feature_dim=None, use_projection=True, **kwargs):
+    """LSNet-B for Artist Style Classification"""
     model = _create_lsnet_artist(
         "lsnet_b_artist",
         pretrained=pretrained,
@@ -211,6 +238,7 @@ def lsnet_b_artist(num_classes=1000, distillation=False, pretrained=False,
 @register_model
 def lsnet_l_artist(num_classes=1000, distillation=False, pretrained=False,
                    feature_dim=None, use_projection=True, **kwargs):
+    """LSNet-L for Artist Style Classification (Large model for massive training)"""
     model = _create_lsnet_artist(
         "lsnet_l_artist",
         pretrained=pretrained,
@@ -218,9 +246,9 @@ def lsnet_l_artist(num_classes=1000, distillation=False, pretrained=False,
         distillation=distillation,
         img_size=224,
         patch_size=8,
-        embed_dim=[160, 320, 480, 640],
-        depth=[6, 8, 12, 14],
-        num_heads=[4, 4, 4, 4],
+        embed_dim=[160, 320, 480, 640],  # 更大的embed_dim
+        depth=[6, 8, 12, 14],           # 更深的网络
+        num_heads=[4, 4, 4, 4],          # 更多的注意力头
         feature_dim=feature_dim,
         use_projection=use_projection,
         **kwargs
@@ -231,6 +259,7 @@ def lsnet_l_artist(num_classes=1000, distillation=False, pretrained=False,
 @register_model
 def lsnet_xl_artist(num_classes=1000, distillation=False, pretrained=False,
                     feature_dim=None, use_projection=True, **kwargs):
+    """LSNet-XL for Artist Style Classification (Extra Large model for massive datasets with 100k+ classes)"""
     model = _create_lsnet_artist(
         "lsnet_xl_artist",
         pretrained=pretrained,
@@ -238,11 +267,32 @@ def lsnet_xl_artist(num_classes=1000, distillation=False, pretrained=False,
         distillation=distillation,
         img_size=224,
         patch_size=8,
-        embed_dim=[192, 384, 576, 768],
-        depth=[8, 12, 16, 20],
-        num_heads=[6, 6, 6, 6],
+        embed_dim=[192, 384, 576, 768],  # 超大embed_dim，支持10万+类别
+        depth=[8, 12, 16, 20],           # 超深网络，学习复杂特征
+        num_heads=[6, 6, 6, 6],           # 更多注意力头
         feature_dim=feature_dim,
         use_projection=use_projection,
         **kwargs
     )
     return model
+
+
+@register_model
+def lsnet_xl_artist_448(num_classes=50000, distillation=False, pretrained=False,
+                        feature_dim=None, use_projection=True, **kwargs):
+    """LSNet-XL-448 for Artist Style Classification (Extra Large model with 448x448 input for massive datasets with 50k+ classes)"""
+    model = _create_lsnet_artist(
+        "lsnet_xl_artist_448",
+        pretrained=pretrained,
+        num_classes=num_classes, 
+        distillation=distillation,
+        img_size=448,
+        patch_size=8,
+        embed_dim=[192, 384, 576, 768],  # 超大embed_dim，支持10万+类别
+        depth=[8, 12, 16, 20],           # 超深网络，学习复杂特征
+        num_heads=[6, 6, 6, 6],           # 更多注意力头
+        feature_dim=feature_dim,
+        use_projection=use_projection,
+        **kwargs
+    )
+    return model