设计一个多层感知机mlp网络。 我希望看到的输出为：Flatten output shape torchSize1 1024Linear output shape torchSize1 2048 Linear weights mean tensor08631 Linear biass mean tensor0Sigmoid output shape torchSize1 2

import torch import torch.nn as nn

class MLP(nn.Module): def init(self): super(MLP, self).init() self.flatten = nn.Flatten() self.linear1 = nn.Linear(32323, 2048) self.sigmoid = nn.Sigmoid() self.linear2 = nn.Linear(2048, 512) self.relu = nn.ReLU() self.linear3 = nn.Linear(512, 20)

def forward(self, x):
    x = self.flatten(x)
    x = self.linear1(x)
    print("Linear output shape: \t", x.shape)
    print("\t Linear weight's mean: \t", torch.mean(self.linear1.weight))
    print("\t Linear bias's mean: \t", torch.mean(self.linear1.bias))
    x = self.sigmoid(x)
    print("Sigmoid output shape: \t", x.shape)
    x = self.linear2(x)
    print("Linear output shape: \t", x.shape)
    print("\t Linear weight's mean: \t", torch.mean(self.linear2.weight))
    print("\t Linear bias's mean: \t", torch.mean(self.linear2.bias))
    x = self.relu(x)
    print("ReLU output shape: \t", x.shape)
    x = self.linear3(x)
    print("Linear output shape: \t", x.shape)
    print("\t Linear weight's mean: \t", torch.mean(self.linear3.weight))
    print("\t Linear bias's mean: \t", torch.mean(self.linear3.bias))
    return x

model = MLP() x = torch.randn(1, 3, 32, 32) output = model(x) print("Flatten output shape: \t", output.shape