代码之家 › 专栏 › 技术社区 › Kevin

ConvNeXt torchvision-指定输入通道

torchvision torch pytorch

0

Kevin · 技术社区 · 3 年前

如何更改torchvision中的输入通道数 ConvNeXt 模型我正在处理灰度图像,想要1个输入通道而不是3个。

import torch
from torchvision.models.convnext import ConvNeXt, CNBlockConfig

# this is the given configuration for the 'tiny' model
block_setting = [
    CNBlockConfig(96, 192, 3),
    CNBlockConfig(192, 384, 3),
    CNBlockConfig(384, 768, 9),
    CNBlockConfig(768, None, 3),
]

model = ConvNeXt(block_setting)

# my sample image (N, C, W, H) = (16, 1, 50, 50)
im = torch.randn(16, 1, 50, 50)
# forward pass
model(im)

输出:

RuntimeError: Given groups=1, weight of size [96, 3, 4, 4], expected input[16, 1, 50, 50] to have 3 channels, but got 1 channels instead

但是,如果我将输入形状更改为 (16, 3, 50, 50) 它似乎运行良好。

torchvision源代码似乎基于 github implementation 但是我具体在哪里 in_chans 使用torchvision接口?

0 回复 | 直到 3 年前

1

3

ndrwnaguib Nikkolai Fernandez 3 年前

您可以重写整个输入层, model._modules["features"][0][0] 是

nn.Conv2d(3, 96, kernel_size=(4, 4), stride=(4, 4))

然后,您只需要更改 in_channels

>>> model._modules["features"][0][0] = nn.Conv2d(1, 96, kernel_size=(4, 4), stride=(4, 4))
>>> model(im)
tensor([[-0.4854, -0.1925,  0.1051,  ..., -0.2310, -0.8830, -0.0251],
        [ 0.3332, -0.4205, -0.3007,  ...,  0.8530,  0.1429, -0.3819],
        [ 0.1794, -0.7546, -0.7835,  ..., -0.8072, -0.0972,  0.7413],
        ...,
        [ 0.1356,  0.0868,  0.6135,  ..., -0.1382, -0.2001,  0.2415],
        [-0.1612, -0.4812,  0.1271,  ..., -0.6594,  0.2706,  1.0833],
        [ 0.0243, -0.5039, -0.4086,  ...,  0.4233,  0.0389,  0.2787]],
       grad_fn=<AddmmBackward0>)