import torch
from torch import nn


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(100 * 100 * 3, 784),
            nn.PReLU(),
            nn.Linear(784, 256),
            nn.PReLU(),
            nn.Linear(256, 128),
            nn.PReLU(),
            nn.Linear(128, 64),
            nn.PReLU(),
            nn.Linear(64, 32),
            nn.PReLU(),
            nn.Linear(32, 3),
            # 损失函数是交叉熵去掉softmax
            # nn.Softmax(dim=1)
        )

    def forward(self, _x):
        return self.fc(_x)


if __name__ == '__main__':
    x = torch.randn(2, 100 * 100 * 3)
    net = Net()
    out = net(x)
    print(out.shape)
    # print(out)
    pass