AI-agent/demo/mlp_mse/net.py

from torch import nn
import torch

class FullyConnectedNet(nn.Module):
    def __init__(self):
        super().__init__()
        # nn.Sequential() 会按顺序执行每一层
        self.net = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10),

            # 输出函数
            nn.Softmax(dim=1)   #NV结构，激活V
        )
    def forward(self, x):
        # forward() 函数定义函数如何从输入流到输出
        # 输入 x 是一批图片，输出是每张图片对应 10 个数字类别的logits
        return self.net(x)



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