文章目录1. 整体结构2. Conv 模块3. C3k2 模块4. C2PSA 模块5. Classify head6. 源码1. 整体结构可以从头创建 from scratch 一个 YOLO26x-cls 分类模型不使用任何官方代码。该分类模型主要包括 4 个大的模块Conv 模块。C3k2 模块。C2PSA 模块。Classify head。使用 Torchview 画出模型的整体结构采用自底向上的方式结构图如下如果将结构图进一步展开可以看到更多的细节。2. Conv 模块Conv 模块包含了 Conv2d, BatchNorm2d 和 SiLU 三部分。其中 Conv2d 的卷积核有的使用 3x3 大小有的是 1x1 大小。3. C3k2 模块C3k2 模块中把 C3k 模块重复了 2 次。在 Torchview 中将模型展开到第三层结构如图。4. C2PSA 模块C2PSA 由 2 个 PSABlock 组成PSA 模块则只有 1 个 PSABlock如图。5. Classify head在推理模式下模型有 2 个输出结构如下图。但在训练模式下则只输出一个未经 softmax 转化的张量。6. 源码创建模型的完整源码 my_yolo26_cls.py 如下#!/usr/bin/env python3本模块用于实现一个自定义的 YOLO26x-cls 分类模型。 版本号 1.0 日期 2026-06-26 作者 drin201312163.com fromcollections.abcimportSequenceimporttorchfromtorchimportnnclassBaseModel(nn.Module):一个基础模型预先设置好了前向传播方法 forward 作为后续其它模型的基类。 Attributes: model : 一个 nn.Module 模型也可以是一个 nn.ModuleList 。 def__init__(self,model:nn.Module|nn.ModuleList):super().__init__()self.modelmodeldefforward(self,x:torch.Tensor|Sequence[torch.Tensor])-torch.Tensor|Sequence[torch.Tensor]:ifisinstance(self.model,nn.ModuleList):foreach_modelinself.model:ifeach_modelisnotNone:xeach_model(x)else:xself.model(x)returnxclassConv(BaseModel):YOLO26 的卷积模块。包括 3 部分卷积BatchNorm2d 第三部分可能是 SiLU 或 Identity 。def__init__(self,in_channels:int,out_channels:int,stride:int1,kernel_size:int3,groups:int1,use_identity:boolFalse):conv_bn_activationnn.ModuleList([nn.Conv2d(in_channelsin_channels,out_channelsout_channels,kernel_sizekernel_size,stridestride,groupsgroups,padding(kernel_size-1)//2,# 不使用 dilation 。biasFalse),nn.BatchNorm2d(num_featuresout_channels),])ifuse_identity:conv_bn_activation.append(nn.Identity())else:conv_bn_activation.append(nn.SiLU(inplaceFalse))super().__init__(modelconv_bn_activation)classBottleneck(nn.Module):由 2 个 Conv 模块组成。 Attributes: conv_1: 第一个 Conv 模块。 conv_2: 第二个 Conv 模块。 def__init__(self,in_channels:int):初始化部分。 Arguments: in_channels: 输入通道数量。 super().__init__()self.conv_1Conv(in_channelsin_channels,out_channelsin_channels,stride1)self.conv_2Conv(in_channelsin_channels,out_channelsin_channels,stride1)defforward(self,x:torch.Tensor)-torch.Tensor:convedself.conv_1(x)convedself.conv_2(conved)returnxconvedclassC3k(nn.Module):CSP Bottleneck 模块可以修改卷积核 k 的大小。 Attributes: entrance_conv: 入口位置的 Conv 模块。 bottleneck_branch: 一个 ModuleList 是 Bottleneck 分支由一个 Conv 模块和 2 个 Bottleneck 组成。 exit_conv: 出口位置的 Conv 模块。 def__init__(self,in_channels:int,out_channels:int):初始化部分。 Arguments: in_channels: 输入通道数量。 out_channels: 输出通道数量。 super().__init__()self.entrance_convConv(in_channelsin_channels,out_channelsin_channels//2,kernel_size1,stride1)bottleneck_convConv(in_channelsin_channels,out_channelsin_channels//2,kernel_size1,stride1)two_bottlenecksnn.Sequential(Bottleneck(in_channelsin_channels//2),# 使用 Conv 模块默认的卷积核大小 k3 。Bottleneck(in_channelsin_channels//2),)self.bottleneck_branchnn.ModuleList([bottleneck_conv,two_bottlenecks])self.exit_convConv(in_channelsin_channels,out_channelsout_channels,kernel_size1,stride1)defforward(self,x:torch.Tensor)-torch.Tensor:bottleneck_branchx# 该步骤只是为了 Torchview 画图时和 YOLO26 官方代码的生成的图一致。foreachinself.bottleneck_branch:bottleneck_brancheach(bottleneck_branch)entrance_convself.entrance_conv(x)# 根据 YOLO26 源码C3k 中先拼接 bottleneck 再拼接 Conv 部分。xtorch.cat((bottleneck_branch,entrance_conv),dim1)returnself.exit_conv(x)classC3k2(nn.Module):把 C3k 模块重复 2 次。 Attributes: entrance_conv: 入口位置的 Conv 模块。 c3k_1: 第 1 个 C3k 模块。 c3k_2: 第 2 个 C3k 模块。 exit_conv: 出口位置的 Conv 模块。 def__init__(self,in_channels:int,out_channels:int):初始化部分。 Arguments: in_channels: 输入通道数量。 out_channels: 输出通道数量。 super().__init__()# 源码中 C3k2 的 2 个 Conv 均使用 kernel 1 。self.entrance_convConv(in_channelsin_channels,out_channelsin_channels,kernel_size1,stride1)self.c3k_1C3k(in_channelsin_channels//2,out_channelsin_channels//2)self.c3k_2C3k(in_channelsin_channels//2,out_channelsin_channels//2)self.exit_convConv(in_channelsin_channels*2,out_channelsout_channels,kernel_size1,stride1)defforward(self,x:torch.Tensor)-torch.Tensor:xself.entrance_conv(x)chunkstorch.chunk(x,chunks2,dim1)c3k_1self.c3k_1(chunks[-1])# YOLO26 源码使用 chunk 的最后一个分支输入给 c3kc3k_2self.c3k_2(c3k_1)# YOLO26 源码中拼接顺序是先拼接 chunk 的部分再拼接 2 个 C3k 的部分。xtorch.cat((*chunks,c3k_1,c3k_2),dim1)returnself.exit_conv(x)classAttention(nn.Module):Self attention 模块 。 Attributes: num_heads: 一个整数是注意力模块的 head 数量。 v_channels: 一个整数是 qkv 中的 v 部分的通道数量。 qk_channels: 一个整数是 qkv 中的 q 部分和 k 部分的通道数量。 scale: 一个浮点数是注意力公式中的分母 sqrt(d) 但是把它转换成了乘法的形式即使用时直接乘以 scale 。 qkv: 入口的 Conv 模块用于计算得到 q, k, v 。 v_conv: 用于计算 v 分支的 Conv 模块。 exit_conv: 模块结尾的 Conv 模块。 def__init__(self,in_channels:int,num_heads:int6):初始化部分。 Arguments: in_channels: 输入通道数量。 num_heads: 注意力模块的 head 数量。 super().__init__()self.num_headsnum_heads# 6self.v_channelsin_channels//num_heads# 即源码的 head_dim 等于 64, in_channels 384channel_ratio0.5# 即源码的 attn_ratioself.qk_channelsint(self.v_channels*channel_ratio)# 即源码的 key_dim等于 32 。self.scaleself.qk_channels**-0.5# attention 中要除以系数 sqrt(d)如果写成乘法形式就可以写成 ** -0.5 。qk_channels_num2*num_heads*self.qk_channels# q, k 需要的通道数量相同所以总数乘以 2。qkv_channels_numqk_channels_numin_channels# qkv 总共需要的通道数量v 需要的通道数量就等于为 in_channels 。self.qkvConv(in_channelsin_channels,out_channelsqkv_channels_num,kernel_size1,use_identityTrue,stride1)# 注意 3 个 Conv 中首尾的 kernel 都是 1只有中间 Conv 的 kernel 为 3 并且使用了 group。self.v_convConv(in_channelsin_channels,out_channelsin_channels,kernel_size3,groupsin_channels,use_identityTrue,stride1)self.exit_convConv(in_channelsin_channels,out_channelsin_channels,kernel_size1,use_identityTrue,stride1)defforward(self,x:torch.Tensor)-torch.Tensor:B,C,H,Wx.shape xself.qkv(x)xx.view(x.shape[0],self.num_heads,x.shape[1]//self.num_heads,-1)q,k,vtorch.split(x,split_size_or_sections[self.qk_channels,self.qk_channels,self.v_channels],dim2)attn(q.transpose(2,3) k)*self.scale# 形状为 batch, 6 ,100, 100attnattn.softmax(dim-1)# 概率值在第 3 维度attnattn.transpose(2,3)# 把概率值转置到第 2 维度。attnv attn# 矩阵乘法会把第 2 维度的概率值和 v 进行逐项相加求和即对 v 求加权和。attnattn.view(B,C,H,W)vself.v_conv(v.reshape(B,C,H,W))attnattnvreturnself.exit_conv(attn)classPSABlock(nn.Module):Position-Sensitive Attention block 。 Attributes: attention: 一个 Attention 注意力模块。 expand_shrink: 一个 Sequential由 2 个 Conv 模块组成特征通道数量先放大一倍然后再恢复原样。 def__init__(self,in_channels:int):初始化部分。 Arguments: in_channels: 输入通道数量。 super().__init__()self.attentionAttention(in_channelsin_channels)# 下面第二个 Conv 是一个特例使用 Identity 代替 SiLU 。self.expand_shrinknn.Sequential(Conv(in_channelsin_channels,out_channelsin_channels*2,kernel_size1,stride1),Conv(in_channelsin_channels*2,out_channelsin_channels,kernel_size1,stride1,use_identityTrue),)defforward(self,x:torch.Tensor)-torch.Tensor:xxself.attention(x)returnxself.expand_shrink(x)classC2PSA(nn.Module):包含 2 个 PSABlock 。 Attributes: entrance_conv: 入口位置的 Conv 模块。 channels: 一个列表存放了 2 个特征通道的数量。 two_psablocks: 一个 Sequential包含 2 个 PSABlock 模块。 exit_conv: 出口位置的 Conv 模块。 def__init__(self,in_channels:int):初始化部分。 Arguments: in_channels: 模块的输入通道数量。 super().__init__()self.entrance_convConv(in_channelsin_channels,out_channelsin_channels,kernel_size1,stride1)# YOLO26 源码中用 split 分成同样大小的 2 组所以下面列表乘以 2 。self.channels[in_channels//2]*2self.two_psablocksnn.Sequential(*(PSABlock(in_channelsin_channels//2)for_inrange(2)))self.exit_convConv(in_channelsin_channels,out_channelsin_channels,kernel_size1,stride1)defforward(self,x:torch.Tensor)-torch.Tensor:xself.entrance_conv(x)a,btorch.split(x,split_size_or_sectionsself.channels,dim1)psablocksself.two_psablocks(b)xtorch.cat((a,psablocks),dim1)# 根据源码先拼接 split[0] 。returnself.exit_conv(x)classClassify(nn.Module):分类模型的输出 head 部分。 Attributes: model: 一个 ModuleList包含了分类模型在 head 部分的所有操作。 def__init__(self,in_channels:int,classes_quantity:int):初始化部分。 Arguments: in_channels: 模块的输入通道数量。 classes_quantity: 最终需要识别的类别数量。 super().__init__()out_channels1280# x 模型固定为 1280self.modelnn.ModuleList([Conv(in_channelsin_channels,out_channelsout_channels,kernel_size1,stride1),nn.AdaptiveAvgPool2d(1),# 把特征图变为 1 x 1nn.Flatten(),nn.Dropout(p0.0,inplaceTrue),# 可以去掉该层。nn.Linear(in_featuresout_channels,out_featuresclasses_quantity),])defforward(self,x:torch.Tensor)-torch.Tensor:foreachinself.model:xeach(x)ifself.training:outputxelse:# 推理模式下还要输出概率值形成一个元祖。probabilitytorch.softmax(x,dim1)outputprobability,x# 源码是先输出概率值再输出 logitreturnoutputclassMyYOLO26CLS(BaseModel):手动实现 YOLO26x-cls 分类模型。使用 My 前缀表示是手动实现以便和官方模型进行区别。def__init__(self,task:str,model_size:str):初始化部分。 Arguments: task: 任务的名字如 classify 等。 model_size: 模型的规模如 x, l, n 等。 allowed_tasksclassify,iftasknotinallowed_tasks:raiseValueError(fAllowed tasks:{allowed_tasks}, but got{task}.)allowed_sizesx,# 其它大小的模型待补充。ifmodel_sizenotinallowed_sizes:raiseValueError(fAllowed sizes:{allowed_sizes}, but got{model_size}.)yolo26nn.ModuleList([Conv(in_channels3,out_channels96,stride2),Conv(in_channels96,out_channels192,stride2),C3k2(in_channels192,out_channels192*2),Conv(in_channels384,out_channels384,stride2),C3k2(in_channels384,out_channels384*2),Conv(in_channels768,out_channels768,stride2),C3k2(in_channels768,out_channels768),Conv(in_channels768,out_channels768,stride2),C3k2(in_channels768,out_channels768),C2PSA(in_channels768),])iftaskclassify:yolo26.append(Classify(in_channels768,classes_quantity1000))else:# 其它任务待补充。raiseValueError(Only classify task is supported.)super().__init__(modelyolo26)defmain():my_yolo26_clsMyYOLO26CLS(taskclassify,model_sizex)print(my_yolo26_cls)if__name____main__:main()—————————— 本文结束 —————————— 手动创建 YOLO26 分类模型
文章目录1. 整体结构2. Conv 模块3. C3k2 模块4. C2PSA 模块5. Classify head6. 源码1. 整体结构可以从头创建 from scratch 一个 YOLO26x-cls 分类模型不使用任何官方代码。该分类模型主要包括 4 个大的模块Conv 模块。C3k2 模块。C2PSA 模块。Classify head。使用 Torchview 画出模型的整体结构采用自底向上的方式结构图如下如果将结构图进一步展开可以看到更多的细节。2. Conv 模块Conv 模块包含了 Conv2d, BatchNorm2d 和 SiLU 三部分。其中 Conv2d 的卷积核有的使用 3x3 大小有的是 1x1 大小。3. C3k2 模块C3k2 模块中把 C3k 模块重复了 2 次。在 Torchview 中将模型展开到第三层结构如图。4. C2PSA 模块C2PSA 由 2 个 PSABlock 组成PSA 模块则只有 1 个 PSABlock如图。5. Classify head在推理模式下模型有 2 个输出结构如下图。但在训练模式下则只输出一个未经 softmax 转化的张量。6. 源码创建模型的完整源码 my_yolo26_cls.py 如下#!/usr/bin/env python3本模块用于实现一个自定义的 YOLO26x-cls 分类模型。 版本号 1.0 日期 2026-06-26 作者 drin201312163.com fromcollections.abcimportSequenceimporttorchfromtorchimportnnclassBaseModel(nn.Module):一个基础模型预先设置好了前向传播方法 forward 作为后续其它模型的基类。 Attributes: model : 一个 nn.Module 模型也可以是一个 nn.ModuleList 。 def__init__(self,model:nn.Module|nn.ModuleList):super().__init__()self.modelmodeldefforward(self,x:torch.Tensor|Sequence[torch.Tensor])-torch.Tensor|Sequence[torch.Tensor]:ifisinstance(self.model,nn.ModuleList):foreach_modelinself.model:ifeach_modelisnotNone:xeach_model(x)else:xself.model(x)returnxclassConv(BaseModel):YOLO26 的卷积模块。包括 3 部分卷积BatchNorm2d 第三部分可能是 SiLU 或 Identity 。def__init__(self,in_channels:int,out_channels:int,stride:int1,kernel_size:int3,groups:int1,use_identity:boolFalse):conv_bn_activationnn.ModuleList([nn.Conv2d(in_channelsin_channels,out_channelsout_channels,kernel_sizekernel_size,stridestride,groupsgroups,padding(kernel_size-1)//2,# 不使用 dilation 。biasFalse),nn.BatchNorm2d(num_featuresout_channels),])ifuse_identity:conv_bn_activation.append(nn.Identity())else:conv_bn_activation.append(nn.SiLU(inplaceFalse))super().__init__(modelconv_bn_activation)classBottleneck(nn.Module):由 2 个 Conv 模块组成。 Attributes: conv_1: 第一个 Conv 模块。 conv_2: 第二个 Conv 模块。 def__init__(self,in_channels:int):初始化部分。 Arguments: in_channels: 输入通道数量。 super().__init__()self.conv_1Conv(in_channelsin_channels,out_channelsin_channels,stride1)self.conv_2Conv(in_channelsin_channels,out_channelsin_channels,stride1)defforward(self,x:torch.Tensor)-torch.Tensor:convedself.conv_1(x)convedself.conv_2(conved)returnxconvedclassC3k(nn.Module):CSP Bottleneck 模块可以修改卷积核 k 的大小。 Attributes: entrance_conv: 入口位置的 Conv 模块。 bottleneck_branch: 一个 ModuleList 是 Bottleneck 分支由一个 Conv 模块和 2 个 Bottleneck 组成。 exit_conv: 出口位置的 Conv 模块。 def__init__(self,in_channels:int,out_channels:int):初始化部分。 Arguments: in_channels: 输入通道数量。 out_channels: 输出通道数量。 super().__init__()self.entrance_convConv(in_channelsin_channels,out_channelsin_channels//2,kernel_size1,stride1)bottleneck_convConv(in_channelsin_channels,out_channelsin_channels//2,kernel_size1,stride1)two_bottlenecksnn.Sequential(Bottleneck(in_channelsin_channels//2),# 使用 Conv 模块默认的卷积核大小 k3 。Bottleneck(in_channelsin_channels//2),)self.bottleneck_branchnn.ModuleList([bottleneck_conv,two_bottlenecks])self.exit_convConv(in_channelsin_channels,out_channelsout_channels,kernel_size1,stride1)defforward(self,x:torch.Tensor)-torch.Tensor:bottleneck_branchx# 该步骤只是为了 Torchview 画图时和 YOLO26 官方代码的生成的图一致。foreachinself.bottleneck_branch:bottleneck_brancheach(bottleneck_branch)entrance_convself.entrance_conv(x)# 根据 YOLO26 源码C3k 中先拼接 bottleneck 再拼接 Conv 部分。xtorch.cat((bottleneck_branch,entrance_conv),dim1)returnself.exit_conv(x)classC3k2(nn.Module):把 C3k 模块重复 2 次。 Attributes: entrance_conv: 入口位置的 Conv 模块。 c3k_1: 第 1 个 C3k 模块。 c3k_2: 第 2 个 C3k 模块。 exit_conv: 出口位置的 Conv 模块。 def__init__(self,in_channels:int,out_channels:int):初始化部分。 Arguments: in_channels: 输入通道数量。 out_channels: 输出通道数量。 super().__init__()# 源码中 C3k2 的 2 个 Conv 均使用 kernel 1 。self.entrance_convConv(in_channelsin_channels,out_channelsin_channels,kernel_size1,stride1)self.c3k_1C3k(in_channelsin_channels//2,out_channelsin_channels//2)self.c3k_2C3k(in_channelsin_channels//2,out_channelsin_channels//2)self.exit_convConv(in_channelsin_channels*2,out_channelsout_channels,kernel_size1,stride1)defforward(self,x:torch.Tensor)-torch.Tensor:xself.entrance_conv(x)chunkstorch.chunk(x,chunks2,dim1)c3k_1self.c3k_1(chunks[-1])# YOLO26 源码使用 chunk 的最后一个分支输入给 c3kc3k_2self.c3k_2(c3k_1)# YOLO26 源码中拼接顺序是先拼接 chunk 的部分再拼接 2 个 C3k 的部分。xtorch.cat((*chunks,c3k_1,c3k_2),dim1)returnself.exit_conv(x)classAttention(nn.Module):Self attention 模块 。 Attributes: num_heads: 一个整数是注意力模块的 head 数量。 v_channels: 一个整数是 qkv 中的 v 部分的通道数量。 qk_channels: 一个整数是 qkv 中的 q 部分和 k 部分的通道数量。 scale: 一个浮点数是注意力公式中的分母 sqrt(d) 但是把它转换成了乘法的形式即使用时直接乘以 scale 。 qkv: 入口的 Conv 模块用于计算得到 q, k, v 。 v_conv: 用于计算 v 分支的 Conv 模块。 exit_conv: 模块结尾的 Conv 模块。 def__init__(self,in_channels:int,num_heads:int6):初始化部分。 Arguments: in_channels: 输入通道数量。 num_heads: 注意力模块的 head 数量。 super().__init__()self.num_headsnum_heads# 6self.v_channelsin_channels//num_heads# 即源码的 head_dim 等于 64, in_channels 384channel_ratio0.5# 即源码的 attn_ratioself.qk_channelsint(self.v_channels*channel_ratio)# 即源码的 key_dim等于 32 。self.scaleself.qk_channels**-0.5# attention 中要除以系数 sqrt(d)如果写成乘法形式就可以写成 ** -0.5 。qk_channels_num2*num_heads*self.qk_channels# q, k 需要的通道数量相同所以总数乘以 2。qkv_channels_numqk_channels_numin_channels# qkv 总共需要的通道数量v 需要的通道数量就等于为 in_channels 。self.qkvConv(in_channelsin_channels,out_channelsqkv_channels_num,kernel_size1,use_identityTrue,stride1)# 注意 3 个 Conv 中首尾的 kernel 都是 1只有中间 Conv 的 kernel 为 3 并且使用了 group。self.v_convConv(in_channelsin_channels,out_channelsin_channels,kernel_size3,groupsin_channels,use_identityTrue,stride1)self.exit_convConv(in_channelsin_channels,out_channelsin_channels,kernel_size1,use_identityTrue,stride1)defforward(self,x:torch.Tensor)-torch.Tensor:B,C,H,Wx.shape xself.qkv(x)xx.view(x.shape[0],self.num_heads,x.shape[1]//self.num_heads,-1)q,k,vtorch.split(x,split_size_or_sections[self.qk_channels,self.qk_channels,self.v_channels],dim2)attn(q.transpose(2,3) k)*self.scale# 形状为 batch, 6 ,100, 100attnattn.softmax(dim-1)# 概率值在第 3 维度attnattn.transpose(2,3)# 把概率值转置到第 2 维度。attnv attn# 矩阵乘法会把第 2 维度的概率值和 v 进行逐项相加求和即对 v 求加权和。attnattn.view(B,C,H,W)vself.v_conv(v.reshape(B,C,H,W))attnattnvreturnself.exit_conv(attn)classPSABlock(nn.Module):Position-Sensitive Attention block 。 Attributes: attention: 一个 Attention 注意力模块。 expand_shrink: 一个 Sequential由 2 个 Conv 模块组成特征通道数量先放大一倍然后再恢复原样。 def__init__(self,in_channels:int):初始化部分。 Arguments: in_channels: 输入通道数量。 super().__init__()self.attentionAttention(in_channelsin_channels)# 下面第二个 Conv 是一个特例使用 Identity 代替 SiLU 。self.expand_shrinknn.Sequential(Conv(in_channelsin_channels,out_channelsin_channels*2,kernel_size1,stride1),Conv(in_channelsin_channels*2,out_channelsin_channels,kernel_size1,stride1,use_identityTrue),)defforward(self,x:torch.Tensor)-torch.Tensor:xxself.attention(x)returnxself.expand_shrink(x)classC2PSA(nn.Module):包含 2 个 PSABlock 。 Attributes: entrance_conv: 入口位置的 Conv 模块。 channels: 一个列表存放了 2 个特征通道的数量。 two_psablocks: 一个 Sequential包含 2 个 PSABlock 模块。 exit_conv: 出口位置的 Conv 模块。 def__init__(self,in_channels:int):初始化部分。 Arguments: in_channels: 模块的输入通道数量。 super().__init__()self.entrance_convConv(in_channelsin_channels,out_channelsin_channels,kernel_size1,stride1)# YOLO26 源码中用 split 分成同样大小的 2 组所以下面列表乘以 2 。self.channels[in_channels//2]*2self.two_psablocksnn.Sequential(*(PSABlock(in_channelsin_channels//2)for_inrange(2)))self.exit_convConv(in_channelsin_channels,out_channelsin_channels,kernel_size1,stride1)defforward(self,x:torch.Tensor)-torch.Tensor:xself.entrance_conv(x)a,btorch.split(x,split_size_or_sectionsself.channels,dim1)psablocksself.two_psablocks(b)xtorch.cat((a,psablocks),dim1)# 根据源码先拼接 split[0] 。returnself.exit_conv(x)classClassify(nn.Module):分类模型的输出 head 部分。 Attributes: model: 一个 ModuleList包含了分类模型在 head 部分的所有操作。 def__init__(self,in_channels:int,classes_quantity:int):初始化部分。 Arguments: in_channels: 模块的输入通道数量。 classes_quantity: 最终需要识别的类别数量。 super().__init__()out_channels1280# x 模型固定为 1280self.modelnn.ModuleList([Conv(in_channelsin_channels,out_channelsout_channels,kernel_size1,stride1),nn.AdaptiveAvgPool2d(1),# 把特征图变为 1 x 1nn.Flatten(),nn.Dropout(p0.0,inplaceTrue),# 可以去掉该层。nn.Linear(in_featuresout_channels,out_featuresclasses_quantity),])defforward(self,x:torch.Tensor)-torch.Tensor:foreachinself.model:xeach(x)ifself.training:outputxelse:# 推理模式下还要输出概率值形成一个元祖。probabilitytorch.softmax(x,dim1)outputprobability,x# 源码是先输出概率值再输出 logitreturnoutputclassMyYOLO26CLS(BaseModel):手动实现 YOLO26x-cls 分类模型。使用 My 前缀表示是手动实现以便和官方模型进行区别。def__init__(self,task:str,model_size:str):初始化部分。 Arguments: task: 任务的名字如 classify 等。 model_size: 模型的规模如 x, l, n 等。 allowed_tasksclassify,iftasknotinallowed_tasks:raiseValueError(fAllowed tasks:{allowed_tasks}, but got{task}.)allowed_sizesx,# 其它大小的模型待补充。ifmodel_sizenotinallowed_sizes:raiseValueError(fAllowed sizes:{allowed_sizes}, but got{model_size}.)yolo26nn.ModuleList([Conv(in_channels3,out_channels96,stride2),Conv(in_channels96,out_channels192,stride2),C3k2(in_channels192,out_channels192*2),Conv(in_channels384,out_channels384,stride2),C3k2(in_channels384,out_channels384*2),Conv(in_channels768,out_channels768,stride2),C3k2(in_channels768,out_channels768),Conv(in_channels768,out_channels768,stride2),C3k2(in_channels768,out_channels768),C2PSA(in_channels768),])iftaskclassify:yolo26.append(Classify(in_channels768,classes_quantity1000))else:# 其它任务待补充。raiseValueError(Only classify task is supported.)super().__init__(modelyolo26)defmain():my_yolo26_clsMyYOLO26CLS(taskclassify,model_sizex)print(my_yolo26_cls)if__name____main__:main()—————————— 本文结束 ——————————