import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models

# Define ResNet model
class ResNet(nn.Module):
    def __init__(self, num_classes):
        super(ResNet, self).__init__()
        self.resnet = models.resnet18(pretrained=True)
        self.resnet.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=1, padding=1, bias=False)
        num_ftrs = self.resnet.fc.in_features
        self.resnet.fc = nn.Linear(num_ftrs, num_classes)

如上代码片段,需要在init的时候进行修改,如果通过resnet的一个实例进行修改,在模型使用的时候还是没有变化。

参考了ResNet训练单通道图像分类网络(Pytorch)_望~的博客-CSDN博客_单通道图像resnet不过该博主给出的方法并不work,原因就是他是对实例化的ResNet进行修改成员变量,而不是在init函数中。

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐