model篇version2(TinyModel)

微型Model

为了方法后续测试Train, Trainer等文件,这里引入由MLP多层感知机搭建的微型model,如下:

import torch
import torch.nn as nn

class TinyNet(nn.Module):
    def __init__(self, input=28*28, output=28*28):
        super().__init__()
        # define any number of nn.Modules (or use your current ones)
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))
    def  forward(self, x):
        y = self.encoder(x)
        z = self.decoder(y)
        return z
Net = TinyNet
if __name__=="__main__":
    net = TinyNet()
    x = torch.randn((28*28))
    z = net(x)
    pass

你可能感兴趣的:(yuque,python)