PyTorch TypeError: flatten() takes at most 1 argument (2 given)

Question

我正在尝试运行这个自定义的 PyTorch 程序：

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten(start_dim=1)
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(2142, 51),
            nn.ReLU(),
            nn.Linear(51, 1)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork().to(device)

以上代码是自定义的NeuralNetwork。我在下面定义了训练循环。

def train_loop(data, y, model, loss_fn, optimizer):
    for i in range(data.shape[0]):
        # Compute model prediction and loss
        pred = model(data[i, :, :])
        loss = loss_fn(pred, y[i, :])

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
        print("Loss: {}".format(loss.item()))

以下是我想训练它的方式

final_data = torch.randn(500, 42, 51)
final_output = torch.randn(500, 1)
learning_rate = 1e-3
batch_size = 1
epochs = 5
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

epochs = 10
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(final_data, final_output, model, loss_fn, optimizer)

print('Done!!')

变量final_data的形状为500x42x51。变量 final_output 的形状为 500x1。我一直在尝试运行以上数据 10 个时期，但我总是以这个错误结束：

Epoch 1
-------------------------------
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-17-3fb934698ecf> in <module>()
      9 for t in range(epochs):
     10     print(f"Epoch {t+1}\n-------------------------------")
---> 11     train_loop(final_data, final_output, model, loss_fn, optimizer)
     12 
     13 print('Done!!')

4 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/flatten.py in forward(self, input)
     38 
     39     def forward(self, input: Tensor) -> Tensor:
---> 40         return input.flatten(self.start_dim, self.end_dim)
     41 
     42     def extra_repr(self) -> str:

TypeError: flatten() takes at most 1 argument (2 given)

输出基本上是0或1之间的分类。就 PyTorch 而言，我仍然是新手，希望得到一些帮助来解决这个问题并了解问题所在。

谢谢

Answer 1

错误有点误导，如果您尝试运行来自新内核的代码，问题出在其他地方...

您的代码存在多个问题：

您没有为目标张量使用正确的形状 y，这里它应该有一个维度，因为输出张量是二维的。
目标张量应该是dtype Long
当迭代数据并使用 data[i, :, :]（和 y[i, :]）选择输入（和目标）时，您实际上是在删除批处理轴。然而，所有内置 nn.Module 都使用批处理轴。您可以做一个切片来避免这种副作用：分别使用 data[i:i+1] 和 y[i:i+1]。 另请注意 x[j, :, :] 等同于 x[j]。

也就是说，使用交叉熵损失是不合理的。您正在输出单个 logit，因此使用交叉熵损失没有意义。

您可以在模型的最后一层输出两个 logits，
或切换到另一个损失函数，例如二元交叉熵损失（使用包含 sigmoid 激活的 nn.BCELoss or a nn.BCEWithLogitsLoss）。

在这种情况下，目标向量应该是 dtype float，它的形状应该等于 pred.

def train_loop(data, y, model, loss_fn, optimizer):
    for i in range(data.shape[0]):
        pred = model(data[i:i+1])
        loss = loss_fn(pred, y[i:i+1].float())
        # [...]

final_data = torch.randn(500, 42, 51)
final_output = torch.randint(0, 2, (500,1))
loss_fn = nn.BCEWithLogitsLoss()

然后以下将起作用：

>>> train_loop(final_data, final_output, model, loss_fn, optimizer)

PyTorch TypeError: flatten() takes at most 1 argument (2 given)

PyTorch TypeError: flatten() takes at most 1 argument (2 given)

python

pytorch