| 30 | # return decode |
| 31 | |
| 32 | class VariationalEncoder(nn.Module): |
| 33 | def __init__(self,latent_dims): |
| 34 | super(VariationalEncoder,self).__init__() |
| 35 | self.linear1 = nn.Linear(256,128) |
| 36 | self.linear2 = nn.Linear(128,64) |
| 37 | self.linear3 = nn.Linear(64,latent_dims) |
| 38 | self.linear4 = nn.Linear(64,latent_dims) |
| 39 | |
| 40 | self.N = torch.distributions.Normal(0,1) |
| 41 | self.N.loc = self.N.loc.cuda() |
| 42 | self.N.scale = self.N.scale.cuda() |
| 43 | self.kl = 0 |
| 44 | |
| 45 | def forward(self,x): |
| 46 | x = F.relu(self.linear1(x)) |
| 47 | x = F.relu(self.linear2(x)) |
| 48 | mu = self.linear3(x) |
| 49 | sigma = torch.exp(self.linear4(x)) |
| 50 | z = mu + sigma*self.N.sample(mu.shape) |
| 51 | self.kl = (sigma**2 + mu**2 - torch.log(sigma) - 1/2).sum() |
| 52 | return z |
| 53 | class Decoder(nn.Module): |
| 54 | def __init__(self,latent_dims): |
| 55 | super(Decoder,self).__init__() |