CNN模型定义输入和输出随RGB的变化

时间:2020-07-05 17:01:15

标签: machine-learning dataset pytorch cnn

我正在尝试使用CNN来处理交通标志数据集,我已经完成了灰度转换,但想了解使用RGB的影响。我使用的尺寸为32x32x1,但现在为32x32x3。

class ConvNet(nn.Module):
def __init__(self):
    super(ConvNet, self).__init__() 
    self.conv1 = nn.Conv2d(in_channels=1,
                           out_channels=6,
                           kernel_size=(5,5),
                           stride=(1,1))
    self.pool2 = nn.MaxPool2d(kernel_size = (2,2))
    self.conv3 = nn.Conv2d (in_channels=6,
                            out_channels=16,
                            kernel_size=(5,5),
                            stride=(1,1))
    self.pool4 = nn.MaxPool2d(kernel_size=(2,2))
    self.conv5 = nn.Conv2d (in_channels=16,
                           out_channels=120,
                           kernel_size=(5,5),
                           stride=(1,1))
    self.fc6 = nn.Linear(in_features=120,
                         out_features=84)
    self.fc7 = nn.Linear(in_features=84,
                         out_features=62)
def forward(self,x):
    x = x.view([-1, 1, 32, 32])
    x = nn.functional.relu(self.conv1(x))
    x = self.pool2(x)
    x = nn.functional.relu(self.conv3(x))
    x = self.pool4(x)
    x = nn.functional.relu(self.conv5(x))
    x = x.view(-1, 120)
    x = nn.functional.relu(self.fc6(x))
    x = self.fc7(x)
    return x

我不明白该怎么做,是用附加通道定义层。希望能为您提供帮助。

1 个答案:

答案 0 :(得分:0)

由于仅更改输入层的通道,因此只需要将self.conv1的输入通道更改为3,而不是1:self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=(5,5), stride=(1,1))。这样做,您还必须将展平形状[-1, 1, 32, 32]更改为[-1, 3, 32, 32],因为您现在拥有3个通道。