在Pytorch中创建具有动态隐藏层数量的前馈神经网络模型

import torch
import torch.nn as nn

class DynamicFeedForward(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size):
        super(DynamicFeedForward, self).__init__()
        self.input_size = input_size
        self.hidden_sizes = hidden_sizes
        self.output_size = output_size
        
        self.input_layer = nn.Linear(input_size, hidden_sizes[0])
        self.hidden_layers = nn.ModuleList()
        for i in range(len(hidden_sizes) - 1):
            self.hidden_layers.append(nn.Linear(hidden_sizes[i], hidden_sizes[i+1]))
        self.output_layer = nn.Linear(hidden_sizes[-1], output_size)
        
    def forward(self, x):
        x = torch.relu(self.input_layer(x))
        for hidden_layer in self.hidden_layers:
            x = torch.relu(hidden_layer(x))
        x = self.output_layer(x)
        return x

input_size = 10
hidden_sizes = [20, 30, 40]  # 可以根据需求调整隐藏层的数量和大小
output_size = 5

model = DynamicFeedForward(input_size, hidden_sizes, output_size)

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 准备训练数据
input_data = torch.randn(100, input_size)
target_data = torch.randn(100, output_size)

# 进行训练
for epoch in range(100):
    optimizer.zero_grad()
    output = model(input_data)
    loss = criterion(output, target_data)
    loss.backward()
    optimizer.step()