def __init__(self, dim_hids=128, num_inds=32): super().__init__() self.encoder = nn.Sequential(FixupResUnit(1, 32, stride=2), FixupResUnit(32, 64, stride=2), FixupResUnit(64, dim_hids, stride=2), nn.AdaptiveAvgPool2d(1)) self.mab1 = MAB(dim_hids, dim_hids, dim_hids) self.isab = StackedISAB(dim_hids, dim_hids, num_inds, 4) self.mab2 = MAB(dim_hids, dim_hids, dim_hids) self.fc = nn.Linear(dim_hids, 1)
def __init__(self, dim_lats, dim_hids=128, num_inds=32): super().__init__() self.encoder = nn.Sequential(View(-1, 784), WN(nn.Linear(784, dim_hids)), nn.ELU(), WN(nn.Linear(dim_hids, dim_hids)), nn.ELU(), WN(nn.Linear(dim_hids, dim_hids))) self.isab1 = StackedISAB(dim_hids, dim_hids, num_inds, 4) self.pma = PMA(dim_hids, dim_hids, 1) self.fc1 = nn.Linear(dim_hids, dim_hids) self.posterior = Normal(dim_lats, use_context=True, context_enc=nn.Linear(2 * dim_hids, 2 * dim_lats)) self.prior = FlowDistribution( MAF(dim_lats, dim_hids, 4, dim_context=dim_hids, inv_linear=True), Normal(dim_lats)) self.decoder = nn.Sequential( WN(nn.Linear(dim_lats + dim_hids, dim_hids)), nn.ELU(), WN(nn.Linear(dim_hids, dim_hids)), nn.ELU(), WN(nn.Linear(dim_hids, 784)), View(-1, 1, 28, 28)) self.likel = Bernoulli((1, 28, 28), use_context=True) self.mab = MAB(dim_hids, dim_hids, dim_hids) self.isab2 = StackedISAB(dim_hids, dim_hids, num_inds, 4) self.fc2 = nn.Linear(dim_hids, 1)
def __init__(self, mvn, dim_hids=128, num_inds=32): super().__init__() self.mvn = mvn self.isab1 = StackedISAB(mvn.dim, dim_hids, num_inds, 4) self.pma = PMA(dim_hids, dim_hids, 1) self.fc1 = nn.Linear(dim_hids, mvn.dim_params) self.mab = MAB(dim_hids, dim_hids, dim_hids) self.isab2 = StackedISAB(dim_hids, dim_hids, num_inds, 4) self.fc2 = nn.Linear(dim_hids, 1)
def __init__(self, dim_hids=256, num_inds=32): super().__init__() self.flow = FlowDistribution( MAF(640, dim_hids, 4, dim_context=dim_hids, inv_linear=True), Normal(640, use_context=False)) self.isab1 = StackedISAB(640, dim_hids, num_inds, 4, ln=True, p=0.2) self.pma = PMA(dim_hids, dim_hids, 1) self.fc1 = nn.Linear(dim_hids, dim_hids) nn.init.uniform_(self.fc1.weight, a=-1e-4, b=1e-4) nn.init.constant_(self.fc1.bias, 0.0) self.mab = MAB(dim_hids, dim_hids, dim_hids) self.isab2 = StackedISAB(dim_hids, dim_hids, num_inds, 4) self.fc2 = nn.Linear(dim_hids, 1)
def __init__(self, dim_inputs, dim_hids=128, num_inds=32, dim_context=128, num_blocks=4): super().__init__() self.flow = FlowDistribution( MAF(dim_inputs, dim_hids, num_blocks, dim_context=dim_context), Normal(dim_inputs, use_context=False)) self.isab1 = StackedISAB(dim_inputs, dim_hids, num_inds, 4) self.pma = PMA(dim_hids, dim_hids, 1) self.fc1 = nn.Linear(dim_hids, dim_context) self.mab = MAB(dim_hids, dim_hids, dim_hids) self.isab2 = StackedISAB(dim_hids, dim_hids, num_inds, 4) self.fc2 = nn.Linear(dim_hids, 1)
def __init__(self, dim_hids=256, num_inds=32): super().__init__() self.isab = StackedISAB(640, dim_hids, num_inds, 6, p=0.2, ln=True) self.mab = MAB(dim_hids, dim_hids, dim_hids) self.fc = nn.Linear(dim_hids, 1)