from torch import nn
import torch

class FullyConnectedNet(nn.Module):
    def __init__(self):
        super().__init__()
        # nn.Sequential 会按顺序执行每一层。
        self.layer = nn.Sequential(
            nn.Flatten(),             # [batch, 1, 28, 28] -> [batch, 784]
            nn.Linear(28 * 28, 512),  # 784 个像素点映射到 512 个隐藏特征
            nn.ReLU(),                # 增加非线性表达能力
            nn.Linear(512, 256),      # 继续提取更紧凑的隐藏特征
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10),       # 输出 10 个数字类别的分数
        )

    def forward(self, x):
        # forward 定义“数据如何从输入流到输出”。
        # 输入 x 是一批图片，输出是每张图片对应 10 个数字类别的 logits。
        return self.layer(x)


if __name__ == '__main__':
    data = torch.randn(1,1,28,28)
    net = FullyConnectedNet()
    output = net(data)
    print(output)