import torch.nn as nn

class MLP(nn.Module): def init(self): super(MLP, self).init() self.flatten = nn.Flatten() # 拉伸层 self.fc1 = nn.Linear(1024, 2048) # 第一个全连接层 self.sigmoid = nn.Sigmoid() # sigmoid激活函数 self.fc2 = nn.Linear(2048, 512) # 第二个全连接层 self.relu = nn.ReLU() # ReLU激活函数 self.fc3 = nn.Linear(512, 20) # 输出层

    # 全连接层权重初始化方案
    nn.init.uniform_(self.fc1.weight, 0, 1)
    nn.init.uniform_(self.fc2.weight, 0, 1)
    nn.init.uniform_(self.fc3.weight, 0, 1)
    nn.init.constant_(self.fc1.bias, 0)
    nn.init.constant_(self.fc2.bias, 0)
    nn.init.constant_(self.fc3.bias, 0)

def forward(self, x):
    x = self.flatten(x)
    x = self.fc1(x)
    x = self.sigmoid(x)
    x = self.fc2(x)
    x = self.relu(x)
    x = self.fc3(x)
    return x

测试网络结构和权重初始化

net = MLP() input_data = torch.randn(1, 1, 32, 32) output = net(input_data) print("Flatten output shape: \t", output.shape) print("Linear output shape: \t", net.fc1.weight.shape) print("\t Linear weight's mean: \t", net.fc1.weight.mean()) print("\t Linear bias's mean: \t", net.fc1.bias.mean()) output = net.sigmoid(output) print("Sigmoid output shape: \t", output.shape) print("Linear output shape: \t", net.fc2.weight.shape) print("\t Linear weight's mean: \t", net.fc2.weight.mean()) print("\t Linear bias's mean: \t", net.fc2.bias.mean()) output = net.relu(output) print("ReLU output shape: \t", output.shape) print("Linear output shape: \t", net.fc3.weight.shape) print("\t Linear weight's mean: \t", net.fc3.weight.mean()) print("\t Linear bias's mean: \t", net.fc3.bias.mean()