使用pytorch在神经网络中定义展平层

Define flatten layer in neural network using pytorch

我试图在启动全连接层之前定义一个展平层。由于我的输入是一个形状为 (512, 2, 2) 的张量,所以我想在 FC 层之前展平这个张量。

我曾经遇到过这个错误:

empty(): argument 'size' must be tuple of ints, but found element of type Flatten at pos 2
import torch.nn as nn
class Network(nn.Module):
    def __init__(self):
        super(Network,self).__init__()
        self.flatten=nn.Flatten()
        self.fc1=nn.Linear(self.flatten,512)
        self.fc2=nn.Linear(512,256)
        self.fc3=nn.Linear(256,3)
 
        
    def forward(self,x):
        x=self.flatten(x) # Flatten layer
        x=torch.ReLU(self.fc1(x))  
        x=torch.ReLU(self.fc2(x))
        x=torch.softmax(self.fc3(x))
        return x

这行不正确:

        self.fc1 = nn.Linear(self.flatten, 512)

nn.Linear 的第一个参数 in_features 应该是 int 而不是 nn.Module

在您的情况下,您将 flatten 属性定义为 nn.Flatten 模块:

        self.flatten = nn.Flatten()

要解决此问题,您必须传递 in_features 等于展平后的特征数:

        self.fc1 = nn.Linear(n_features_after_flatten, 512)