如何在CNN中为k FC层使用k通道

Question

我有一个编码器，它输出一个形状为(bn, c * k, 32, 32)的张量。我现在想用shape(bn, k, 1, 2)生成k均值。所以平均值是2维坐标。为此，我希望使用k FC层，而对于每个平均k_I，我只希望使用c通道

因此，我的想法是，我将编码器输出out重塑为一个具有形状(bn, k, c, 32, 32)的5d张量。然后我可以使用展平的out[:, 0]out[:, k]作为k线性层的输入

简单的解决方案是手动定义线性层：

self.fc0 = nn.Linear(c * 32 * 32, 2)
...
self.fck = nn.Linear(c * 32 * 32, 2)

然后我可以为每个平均值定义向前传球，如下所示：

mean_0 = self.fc0(out[:, 0].reshape(bn, -1))
...
mean_k = self.fck(out[:, k].reshape(bn, -1))

有没有更有效的方法

Answer 1

您可以使用nn.ModuleList执行类似操作：

import torch
import torch.nn as nn
import torch.nn.functional as F

class fclist(nn.Module):
    def __init__(self, k):
        super().__init__()
        '''
        k: no. of clusters
        '''
        self.k = k
        '''
        .
        .
        .
        Other previous layers
        .
        .
        '''

        c = 1

        self.out_layers = nn.ModuleList()
        for i in range(k):
            self.out_layers.append(nn.Linear(c*32*32, 2))

    def forward(self, x):
        '''
        .
        .
        .
        pass throgh previous layers
        .
        .
        '''
        x = [layer(x) for layer in self.out_layers]
        return x

样本输出：

>>> net =  fclist(k=3)
>>> inp = torch.randn(1, 1*32*32)
>>> net(inp)
[tensor([[-0.7319, -0.2686]], grad_fn=<AddmmBackward>), tensor([[-0.6248,  0.9180]], grad_fn=<AddmmBackward>), tensor([[0.2532, 0.1387]], grad_fn=<AddmmBackward>)]