Pytorch 中的循环优化

Question

我正在尝试找到一种方法来防止 Pytorch 中出现令人痛苦的缓慢 for 循环。基本上，我有一个张量，我想将它分成几部分并将这些部分输入我的模型，在精神上类似于各种分组卷积。

        self.C = C
        self.block = Block(C, 3, 64)
    def forward(self, x):
        x_shape = x.shape
        x = torch.flatten(x, start_dim=1, end_dim=-1).unsqueeze(1)
        x = torch.split(x, self.C, -1)
        attention = []
        for i in x:
            attended = self.block(i)
            attention.append(attended)
        attention = torch.stack(attention, 1)

我认为，C 的小值加上大张量会使这个操作出奇地慢得多，因为上面的代码运行了 Python for 循环。但是，当我将批处理维度交换为“C”维度并循环遍历批处理维度时，这会导致显着的加速，但是对我来说仍然感觉很笨拙，并且在足够大的批处理大小下可能仍然会很慢。我想要一种方法来解决这个问题，同时仍然保持批处理暗淡完整并避免 for 循环。我想我正在寻找的是一种向我的模型添加第二个批次维度的方法，或者类似的方法。

除了上面描述的稍微有点hacky的方法之外，还有其他方法可以解决这个问题吗？

编辑： MWE：（假装单个线性层就像一个分裂的注意力层......）

import torch

class Net(torch.nn.Module):
    def __init__(self, split_size):
        super().__init__()
        self.split_size = split_size
        self.linear = torch.nn.Linear(split_size, split_size)
    def forward(self, x):
        #Slow implementation:

        #Input is B,C,H and is flattened to B,C*H.
        y = x.flatten(start_dim=1, end_dim=-1)

        y_split = torch.split(y, self.split_size, 1) #Tensor is split and each piece is fed into the model...

        outs = []

        for i in y_split:
            i_out = self.linear(i)
            outs.append(i_out)
        y = torch.cat(outs, 1)
        print(y.shape)

        #Fast implementation using batch dims, but possibly slower for large batches...
        y = x.flatten(start_dim=1, end_dim=-1)
        y_split = torch.split(y, self.split_size, 1)

        y = torch.stack(y_split, 0)
        outs = []
        for i in torch.split(y, 1, 1):
            i_out = self.linear(i.squeeze(1)).unsqueeze(0)
            outs.append(i_out)
        y = torch.cat(outs, 0)
        y = y.flatten(start_dim=1, end_dim=-1)
        print(y.shape)
        return y
        
if __name__ == "__main__":
    net = Net(32)
    net(torch.randn(256,3,32,32))
    net(torch.randn(32,3,32,32))

Answer 1

我认为 2D 输入将在具有相同网络层的每个通道中处理。