ConvNeXt 的 torchvision 版本 模型结构的代码实现flyfish# flyfish convnext_tiny.pyfromcollections.abcimportSequencefromfunctoolsimportpartialfromtypingimportAny,Callable,Optionalimporttorchfromtorchimportnn,Tensorfromtorch.nnimportfunctionalasFfromtorchvision.ops.miscimportConv2dNormActivation,Permutefromtorchvision.ops.stochastic_depthimportStochasticDepth __all__[ConvNeXt,convnext_tiny]classLayerNorm2d(nn.LayerNorm):2D 特征图上的 LayerNorm通道维度在最后时做归一化defforward(self,x:Tensor)-Tensor:xx.permute(0,2,3,1)xF.layer_norm(x,self.normalized_shape,self.weight,self.bias,self.eps)xx.permute(0,3,1,2)returnxclassCNBlock(nn.Module):ConvNeXt 基础残差块深度可分离卷积 两层MLP LayerScale 随机深度def__init__(self,dim,layer_scale:float,stochastic_depth_prob:float,norm_layer:Optional[Callable[...,nn.Module]]None,)-None:super().__init__()ifnorm_layerisNone:norm_layerpartial(nn.LayerNorm,eps1e-6)self.blocknn.Sequential(nn.Conv2d(dim,dim,kernel_size7,padding3,groupsdim,biasTrue),Permute([0,2,3,1]),norm_layer(dim),nn.Linear(in_featuresdim,out_features4*dim,biasTrue),nn.GELU(),nn.Linear(in_features4*dim,out_featuresdim,biasTrue),Permute([0,3,1,2]),)self.layer_scalenn.Parameter(torch.ones(dim,1,1)*layer_scale)self.stochastic_depthStochasticDepth(stochastic_depth_prob,row)defforward(self,input:Tensor)-Tensor:resultself.layer_scale*self.block(input)resultself.stochastic_depth(result)resultinputreturnresultclassCNBlockConfig:每个 Stage 的配置输入通道、输出通道、Block 数量def__init__(self,input_channels:int,out_channels:Optional[int],num_layers:int,)-None:self.input_channelsinput_channels self.out_channelsout_channels self.num_layersnum_layersclassConvNeXt(nn.Module):ConvNeXt 完整网络结构def__init__(self,block_setting:list[CNBlockConfig],stochastic_depth_prob:float0.0,layer_scale:float1e-6,num_classes:int1000,block:Optional[Callable[...,nn.Module]]None,norm_layer:Optional[Callable[...,nn.Module]]None,**kwargs:Any,)-None:super().__init__()ifnotblock_setting:raiseValueError(block_setting 不能为空)elifnot(isinstance(block_setting,Sequence)andall([isinstance(s,CNBlockConfig)forsinblock_setting])):raiseTypeError(block_setting 必须是 CNBlockConfig 列表)ifblockisNone:blockCNBlockifnorm_layerisNone:norm_layerpartial(LayerNorm2d,eps1e-6)layers:list[nn.Module][]# Stem 输入层firstconv_output_channelsblock_setting[0].input_channels layers.append(Conv2dNormActivation(3,firstconv_output_channels,kernel_size4,stride4,padding0,norm_layernorm_layer,activation_layerNone,biasTrue,))total_stage_blockssum(cnf.num_layersforcnfinblock_setting)stage_block_id0forcnfinblock_setting:# 每个 Stage 的 Block 堆叠stage:list[nn.Module][]for_inrange(cnf.num_layers):# 随网络深度线性增加随机深度概率sd_probstochastic_depth_prob*stage_block_id/(total_stage_blocks-1.0)stage.append(block(cnf.input_channels,layer_scale,sd_prob))stage_block_id1layers.append(nn.Sequential(*stage))# 下采样层ifcnf.out_channelsisnotNone:layers.append(nn.Sequential(norm_layer(cnf.input_channels),nn.Conv2d(cnf.input_channels,cnf.out_channels,kernel_size2,stride2),))self.featuresnn.Sequential(*layers)self.avgpoolnn.AdaptiveAvgPool2d(1)lastblockblock_setting[-1]lastconv_output_channels(lastblock.out_channelsiflastblock.out_channelsisnotNoneelselastblock.input_channels)self.classifiernn.Sequential(norm_layer(lastconv_output_channels),nn.Flatten(1),nn.Linear(lastconv_output_channels,num_classes))# 权重初始化forminself.modules():ifisinstance(m,(nn.Conv2d,nn.Linear)):nn.init.trunc_normal_(m.weight,std0.02)ifm.biasisnotNone:nn.init.zeros_(m.bias)defforward(self,x:Tensor)-Tensor:xself.features(x)xself.avgpool(x)xself.classifier(x)returnxdefconvnext_tiny(num_classes:int1000,**kwargs:Any)-ConvNeXt: 构造 ConvNeXt-Tiny 模型 Args: num_classes: 分类类别数默认 1000ImageNet Returns: ConvNeXt 模型实例 block_setting[CNBlockConfig(96,192,3),CNBlockConfig(192,384,3),CNBlockConfig(384,768,9),CNBlockConfig(768,None,3),]stochastic_depth_probkwargs.pop(stochastic_depth_prob,0.1)modelConvNeXt(block_setting,stochastic_depth_probstochastic_depth_prob,num_classesnum_classes,**kwargs)returnmodel基础语法1. 类型注解Typing 体系代码所有函数/类的参数定义如Optional[Callable[..., nn.Module]]、list[CNBlockConfig]、Tensor语法Optional可选类型可为 None、Callable可调用对象/函数、Sequence序列类型、Any任意类型设计作用静态类型检查提前发现参数类型错误提升代码可读性一眼可知入参/返回值的类型与含义是深度学习开源项目的通用工程规范。2.functools.partial偏函数代码norm_layer partial(LayerNorm2d, eps1e-6)语法作用预先固定函数的部分参数生成一个新的“简化版函数”。这里的设计意图LayerNorm 本身需要传入eps参数而整个网络所有归一化层的eps都统一为 1e-6。用partial提前固定该参数后后续调用norm_layer(dim)时只需要传通道数即可避免重复传参也方便统一修改全局配置。这是 PyTorch 中定义归一化/激活层的最常用写法。3. 可变参数与关键字参数透传代码**kwargs: Any、**kwargs传参、kwargs.pop(...)**kwargs接收任意数量的关键字参数在模型构造函数中用于透传额外配置比如自定义初始化、额外正则参数保证接口扩展性。kwargs.pop(stochastic_depth_prob, 0.1)从 kwargs 中取出指定参数同时设置默认值取出后该参数会从 kwargs 中移除避免后续透传时报错。这是工厂函数的标准写法既支持自定义参数又不影响剩余参数向底层传递。PyTorch 自定义网络1. 自定义模块的标准范式继承nn.Module代码LayerNorm2d、CNBlock、ConvNeXt全部继承自nn.Module固定写法规则__init__中定义所有可学习参数/子模块必须先调用super().__init__()重写forward方法实现前向计算逻辑子模块如nn.Conv2d、nn.Sequential赋值为实例属性后会自动被 PyTorch 追踪参数会加入模型的参数列表。这是所有 PyTorch 自定义网络的基础框架。2.nn.Sequential序列容器代码self.block nn.Sequential(...)、self.features nn.Sequential(*layers)、self.classifier语法作用按顺序堆叠多个层/模块自动按顺序执行前向计算无需在forward中手动逐层调用简化代码。适用场景固定顺序的线性数据流是卷积块、分类头、主干网络堆叠的最常用容器。扩展写法nn.Sequential(*stage)用*解包列表把列表里的模块逐个传入 Sequential。3.nn.Parameter可训练参数代码self.layer_scale nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)语法作用把普通张量包装成模型可训练参数会自动被优化器识别并更新梯度。设计细节形状设为(dim, 1, 1)利用 PyTorch 的广播机制可以直接和[B, C, H, W]形状的特征图做逐元素相乘实现「每个通道单独缩放」的效果初始化乘以layer_scale小系数1e-6保证训练初期残差分支权重极低网络近似恒等映射稳定深层训练。4. 模块遍历与权重初始化代码forminself.modules():ifisinstance(m,(nn.Conv2d,nn.Linear)):nn.init.trunc_normal_(m.weight,std0.02)ifm.biasisnotNone:nn.init.zeros_(m.bias)self.modules()递归遍历模型内所有子模块返回所有嵌套的 nn.Module。isinstance判断层类型针对不同类型的层执行不同的初始化策略。nn.init系列函数PyTorch 内置的初始化工具这里用截断正态分布初始化权重、偏置置零是 ConvNeXt/ViT 类模型的标准初始化方式。模块的设计细节1.LayerNorm2d适配 CNN 格式的层归一化代码自定义LayerNorm2d类继承nn.LayerNorm语法重写forward通过两次permute切换维度格式原生nn.LayerNorm默认对最后一维做归一化适合序列格式[B, N, C]CNN 特征图默认格式是[B, C, H, W]通道在第2维无法直接用 LayerNorm先permute(0,2,3,1)转成[B, H, W, C]通道在最后做完归一化后再转回去。这是卷积网络中使用 LayerNorm 的标准工程实现等价于对通道维度做归一化。2.CNBlock残差块的完整实现深度卷积写法nn.Conv2d(dim, dim, kernel_size7, padding3, groupsdim)groupsdim是深度卷积的分组数等于输入通道数每个通道独立卷积参数量远小于标准卷积。维度置换 Linear 等价 1×1 卷积先Permute把通道放到最后一维再用Linear做通道维度的线性变换数学上和 1×1 点卷积完全等价好处是可以直接复用原生 LayerNorm写法更简洁。残差连接实现result input主分支是恒等映射残差分支做特征变换最后逐元素相加是残差网络的标准写法。随机深度的位置StochasticDepth放在残差分支末尾、残差相加之前训练时按概率随机丢弃整个残差分支等效于动态减少网络深度。3. 配置化组网CNBlockConfig 循环构建代码CNBlockConfig配置类 for cnf in block_setting循环生成层设计模式配置驱动网络构建。用配置类定义每个 Stage 的「输入通道、输出通道、Block数量」主类只负责循环读取配置、生成对应层。优势扩展性极强想做不同大小的模型Tiny/Small/Base只需要改配置列表不用改网络主逻辑代码复用度高结构清晰是工业级模型库的通用设计。4. 随机深度概率线性递增代码sd_prob stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)语法逻辑按 Block 的深度序号从 0 到最大概率线性递增。设计意图浅层网络特征基础不需要太多随机丢弃深层网络语义复杂正则强度更高。这是 ConvNeXt/Swin 等深层模型的标准训练策略。工厂函数模式代码convnext_tiny()函数写法逻辑封装好对应模型的配置列表调用主类返回模型实例。作用对用户屏蔽复杂的配置细节用户只需调用convnext_tiny()就能拿到模型不同尺寸的模型对应不同工厂函数符合 PyTorch 官方模型库的接口规范。