行业知识
pytorch修改ConvNeXt-T网络
2026-11-12 07:14:09
来源:互联网转载
使用迁移学习,修改ConvNeXt-T网络,对特征进行融合
import torchimport torch.nn as nnimport torchvision.models as modelsclass CustomConvNeXtT(nn.Module): def __init__(self, in_channels=3, num_classes=2, chunk=2, csv_shape=107, CSV=True): super(CustomConvNeXtT, self).__init__() self.chunk = chunk self.num_classes = num_classes self.CSV = CSV # 加载预训练的ConvNeXt-Tiny模型 convnext = models.convnext_tiny(pretrained=True) # 冻结预训练模型的所有参数 for name, param in convnext.named_parameters(): param.requires_grad = False # 将修改后的模型赋值给自定义的ConvNeXt-T网络 self.model = convnext # 修改第一个卷积层的输入通道数 self.model.features[0][0] = nn.Conv2d(in_channels, 96, kernel_size=4, stride=4) # 获取特征提取器的输出特征维度 num_ftrs = self.model.classifier[2].in_features # 修改分类头部 self.model.classifier = nn.Sequential( nn.LayerNorm(num_ftrs * self.chunk + (csv_shape if CSV else 0), eps=1e-6, elementwise_affine=True), nn.Linear(num_ftrs * self.chunk + (csv_shape if CSV else 0), num_classes) ) def extract_features(self, x): x = self.model.features(x) x = self.model.avgpool(x) x = torch.flatten(x, 1) return x def forward(self, data_DCE, data_T2, csv): data_DCE = self.extract_features(data_DCE) data_T2 = self.extract_features(data_T2) if not self.CSV: csv = torch.ones_like(csv) x = torch.cat((data_DCE, data_T2, csv), dim=1) print(f"Feature size after concatenation: {x.size()}") # 打印特征拼接后的尺寸 output = self.model.classifier(x) return outputif __name__ == '__main__': net = CustomConvNeXtT(in_channels=3, num_classes=2, chunk=2, csv_shape=107, CSV=True) for name, param in net.named_parameters(): print(name, ":", param.requires_grad) data_DCE = torch.randn(64, 3, 224, 224) data_T2 = torch.randn(64, 3, 224, 224) csv = torch.randn(64, 107) output = net(data_DCE, data_T2, csv) print("输出特征尺寸:", output.size())
上一篇:核酸服务器异常怎么处理?
下一篇:新项目怎么部署服务器?