| 8 | from torchinfo import summary |
| 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) |
| 26 | x = self.initial_dense(x) |
| 27 | |
| 28 | z_mean = self.z_mean(x) |
| 29 | z_log_var = self.z_log_var(x) |
| 30 | |
| 31 | return z_mean,z_log_var |
| 32 | |
| 33 | class Decoder(torch.nn.Module): |
| 34 | def __init__(self,hidden_dim = 256,latent_dim = 2,num_features = 784): |
no outgoing calls
no test coverage detected