嘿，这需求很清晰！咱们直接上手实现这个PyTorch神经网络，完全匹配你的要求：2个图像特征输入、1个隐藏层、3个0-1之间的输出（对应正方形、矩形等类别的概率）。

首先，先明确核心需求：

• 输入：2个代表图像特征的浮点数
• 结构：1个带非线性激活的隐藏层
• 输出：3个0-1之间的实数，对应三个类别的概率（比如(1,0,0)代表正方形）

步骤1：导入PyTorch核心库 #

先把需要的工具包导入：

import torch
import torch.nn as nn
import torch.optim as optim

步骤2：定义神经网络类 #

咱们继承PyTorch的nn.Module来构建自定义模型，代码里我会加注释解释每一步：

class ShapeClassifier(nn.Module):
    def __init__(self, hidden_size=16):
        super(ShapeClassifier, self).__init__()
        # 输入层 -> 隐藏层：2个输入特征，hidden_size个隐藏单元
        self.fc1 = nn.Linear(2, hidden_size)
        # 隐藏层的非线性激活函数，这里用ReLU（你也可以换成Tanh等）
        self.relu = nn.ReLU()
        # 隐藏层 -> 输出层：hidden_size个输入，3个输出对应三个类别
        self.fc2 = nn.Linear(hidden_size, 3)
        # 输出层用Softmax，确保输出值在0-1之间且总和为1
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        # 前向传播流程：输入 -> 线性层 -> 激活 -> 线性层 -> Softmax
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.softmax(out)
        return out

小提示：如果后续用CrossEntropyLoss做训练，其实可以把模型里的Softmax去掉——因为CrossEntropyLoss已经内置了LogSoftmax和负对数似然损失，这样数值计算更稳定，避免梯度消失问题。

步骤3：初始化模型 #

选个合适的隐藏层大小（比如16，你可以根据需求调整），实例化模型：

# 初始化模型，隐藏层设为16个单元
model = ShapeClassifier(hidden_size=16)
# 打印模型结构，确认是否符合预期
print(model)

运行后你会看到模型的层级结构，确保输入输出维度正确。

步骤4：测试前向传播 #

用随机张量模拟输入，看看输出是否符合要求：

# 生成1个样本，2个特征的输入张量
test_input = torch.randn(1, 2)
# 得到模型输出
output = model(test_input)

print("模型输出：", output)
print("输出总和：", output.sum().item())  # 应该接近1
print("每个输出值范围：", torch.all(output >= 0) and torch.all(output <= 1))  # 应该返回True

这一步能验证模型的输出确实是0-1之间的实数，且总和为1，完美匹配你的需求。

步骤5：补充训练流程（可选） #

如果要让模型学会从特征预测形状，还需要训练流程。这里给个极简示例：

# 定义损失函数（如果模型没加Softmax，用CrossEntropyLoss；加了的话可以用NLLLoss或BCELoss）
# 这里我们用CrossEntropyLoss，所以先把模型里的Softmax注释掉
criterion = nn.CrossEntropyLoss()
# 定义优化器，用Adam比较常用
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 模拟训练数据：32个样本，每个2个特征；标签是0、1、2对应三个类别
X_train = torch.randn(32, 2)
y_train = torch.randint(0, 3, (32,))

# 训练循环
epochs = 100
for epoch in range(epochs):
    # 清空梯度
    optimizer.zero_grad()
    # 前向传播
    outputs = model(X_train)
    # 计算损失
    loss = criterion(outputs, y_train)
    # 反向传播+更新参数
    loss.backward()
    optimizer.step()
    
    # 每10轮打印一次损失
    if (epoch + 1) % 10 == 0:
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")

训练完成后，模型就能根据输入的图像特征输出对应形状的概率了——你可以取输出中最大值的索引作为预测的类别（比如torch.argmax(output, dim=1)）。

内容的提问来源于stack exchange，提问作者stepanVich