Spaces:
Running
Running
File size: 1,326 Bytes
11c72a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | import torch
import torch.nn as nn
import torch.nn.functional as F
class MLPEncoder(nn.Module):
def __init__(self, vocab_size, num_topic, hidden_dim, dropout):
super().__init__()
self.fc11 = nn.Linear(vocab_size, hidden_dim)
self.fc12 = nn.Linear(hidden_dim, hidden_dim)
self.fc21 = nn.Linear(hidden_dim, num_topic)
self.fc22 = nn.Linear(hidden_dim, num_topic)
self.fc1_drop = nn.Dropout(dropout)
self.z_drop = nn.Dropout(dropout)
self.mean_bn = nn.BatchNorm1d(num_topic, affine=True)
self.mean_bn.weight.requires_grad = False
self.logvar_bn = nn.BatchNorm1d(num_topic, affine=True)
self.logvar_bn.weight.requires_grad = False
def reparameterize(self, mu, logvar):
if self.training:
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
return mu + (eps * std)
else:
return mu
def forward(self, x):
e1 = F.softplus(self.fc11(x))
e1 = F.softplus(self.fc12(e1))
e1 = self.fc1_drop(e1)
mu = self.mean_bn(self.fc21(e1))
logvar = self.logvar_bn(self.fc22(e1))
theta = self.reparameterize(mu, logvar)
theta = F.softmax(theta, dim=1)
theta = self.z_drop(theta)
return theta, mu, logvar
|