From 17453ccbe860a08bce54d93d3120419e0576976e Mon Sep 17 00:00:00 2001 From: James Betker Date: Tue, 17 Aug 2021 09:09:29 -0600 Subject: [PATCH] Revert mods to lrdvae They didn't really change anything --- codes/models/gpt_voice/lucidrains_dvae.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/codes/models/gpt_voice/lucidrains_dvae.py b/codes/models/gpt_voice/lucidrains_dvae.py index 6933a301..cba5c728 100644 --- a/codes/models/gpt_voice/lucidrains_dvae.py +++ b/codes/models/gpt_voice/lucidrains_dvae.py @@ -31,9 +31,9 @@ class ResBlock(nn.Module): super().__init__() self.net = nn.Sequential( conv(chan, chan, 3, padding = 1), - nn.SiLU(), + nn.ReLU(), conv(chan, chan, 3, padding = 1), - nn.SiLU(), + nn.ReLU(), conv(chan, chan, 1) ) @@ -88,14 +88,12 @@ class DiscreteVAE(nn.Module): dec_layers = [] for (enc_in, enc_out), (dec_in, dec_out) in zip(enc_chans_io, dec_chans_io): - for _ in range(num_resnet_blocks): - dec_layers.append(ResBlock(dec_in, conv)) + enc_layers.append(nn.Sequential(conv(enc_in, enc_out, 4, stride = 2, padding = 1), nn.ReLU())) + dec_layers.append(nn.Sequential(conv_transpose(dec_in, dec_out, 4, stride = 2, padding = 1), nn.ReLU())) - enc_layers.append(nn.Sequential(conv(enc_in, enc_out, 4, stride = 2, padding = 1), nn.SiLU())) - dec_layers.append(nn.Sequential(conv_transpose(dec_in, dec_out, 4, stride = 2, padding = 1), nn.SiLU())) - - for _ in range(num_resnet_blocks): - enc_layers.append(ResBlock(enc_out, conv)) + for _ in range(num_resnet_blocks): + dec_layers.insert(0, ResBlock(dec_chans[1], conv)) + enc_layers.append(ResBlock(enc_chans[-1], conv)) if num_resnet_blocks > 0: dec_layers.insert(0, conv(codebook_dim, dec_chans[1], 1)) @@ -204,7 +202,7 @@ if __name__ == '__main__': #v = DiscreteVAE() #o=v(torch.randn(1,3,256,256)) #print(o.shape) - v = DiscreteVAE(channels=1, normalization=None, positional_dims=1, num_tokens=4096, codebook_dim=2048, hidden_dim=256, num_resnet_blocks=2) + v = DiscreteVAE(channels=1, normalization=None, positional_dims=1, num_tokens=4096, codebook_dim=2048, hidden_dim=256) v.eval() o=v(torch.randn(1,1,256)) print(o[-1].shape)