resnet读入灰度图训练和预测

2,446 阅读2分钟

resnet默认输入图像通道为3通道,预训练模型也是3通道的。如果希望输入是单通道灰度图该如何修改呢? 比较简单,只需要修改第一个卷积的输入通道即可。查看resnet的函数初始化可以看到最开始的第一个卷积是3通道的,将其改为单通道。

【1】修改resnet函数定义

# self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)

【2】修改预训练模型的加载,不导入第一个卷积层的权重。 参考:github.com/ShuaiLYU/re…

# if pretrained:
#     state_dict = model_zoo.load_url(model_urls['resnet18'])
#     for key  in list(state_dict.keys() ):
#         if "conv1" in key : del state_dict[key]
#     model.load_state_dict(state_dict, False)

【3】预处理归一化操作 原先读入彩图会有HWC3个维度的数据,3个图像通道分别归一化,同时加一个维度即可有BCHW数据进入tensor,而灰度图只有HW两个维度数据,C通道需要单独添加,归一化数据也需要修改

def load_image(self, image_path):
    # self.RGB_MEAN = np.array([122.67891434, 116.66876762, 104.00698793])
    self.RGB_MEAN = 122.67891434
    # img = cv2.imread(image_path, cv2.IMREAD_COLOR).astype('float32')
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype('float32')
    height, width  = img.shape[:2]
    if height > width:
        img = cv2.copyMakeBorder(img, 0, 0, 0, height - width, cv2.BORDER_CONSTANT, value=(255, 255, 255))
    else:
        img = cv2.copyMakeBorder(img, 0, width - height, 0, 0, cv2.BORDER_CONSTANT, value=(255, 255, 255))
    original_shape = img.shape[:2]
    img = self.resize_image_cn(img)
    img -= self.RGB_MEAN
    img /= 255.
    img = np.expand_dims(img, axis=-1)
    img = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0)
    return img, original_shape


def resize_image_cn(self, img, img_side=1280):
    height, width, _ = img.shape
    if height > img_side:
        new_height = 1280
        new_width = 1280
    else:
        new_height = int(math.ceil(height / 32) * 32)
        new_width = int(math.ceil(new_height / height * width / 32) * 32)
    resized_img = cv2.resize(img, (new_width, new_height))
    return resized_img
    

3通道图像

   def load_image(self, image_path):
        img = cv2.imread(image_path, cv2.IMREAD_COLOR).astype('float32')
        height, width, _ = img.shape
        if height > width:
            img = cv2.copyMakeBorder(img, 0, 0, 0, height - width, cv2.BORDER_CONSTANT, value=(255, 255, 255))
        else:
            img = cv2.copyMakeBorder(img, 0, width - height, 0, 0, cv2.BORDER_CONSTANT, value=(255, 255, 255))
        original_shape = img.shape[:2]
        img = self.resize_image_cn(img)
        img -= self.RGB_MEAN
        img /= 255.
        img = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0)
        return img, original_shape

思考: 初始改为灰度图输入是想减少模型的大小,但实际上将图像改为单通道图输入,仅仅修改了第一个卷积,模型文件的大小基本上并没有太多变化,不会减少很多。反而为了适配这种修改需要修改好几处函数的使用(数据的加载,图像预处理,模型的加载,模型的函数定义),而且再次用到其他场景时如果是彩图场景无法适配,比较麻烦且不太实用,所以并不推荐这种做法,默认加载3通道图像输入就好。 至于如何减小模型,加速推理可以使用比较成熟的量化剪枝等操作。