| 9 | |
| 10 | class Encoder(torch.nn.Module): |
| 11 | def __init__(self,hidden_dim = 512,latent_dim = 2): |
| 12 | super(Encoder, self).__init__() |
| 13 | self.initial_dense = torch.nn.Sequential( |
| 14 | torch.nn.Linear(in_features=784, out_features=hidden_dim), |
| 15 | torch.nn.ReLU(inplace=True), |
| 16 | |
| 17 | torch.nn.Linear(in_features = hidden_dim,out_features=256), |
| 18 | torch.nn.ReLU(inplace=True) |
| 19 | ) |
| 20 | #输出的均值和方差 |
| 21 | self.z_mean = torch.nn.Linear(in_features = 256,out_features=latent_dim) |
| 22 | self.z_log_var = torch.nn.Linear(in_features=256,out_features=latent_dim) |
| 23 | |
| 24 | def forward(self,x): |
| 25 | x = x.view(-1,784) |