微型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