diff --git a/models/criticFull.py b/models/criticFull.py index 438eee1baa4355e148367a75527ed698befa7802..382112db6c4b11eaa0321577de21e9ebe8a8c625 100644 --- a/models/criticFull.py +++ b/models/criticFull.py @@ -5,7 +5,7 @@ import torch.nn.functional as F class CriticEMB(nn.Module): - def __init__(self, isize_1=30, isize_2=48, nc=2, ndf=32, size_embed=16): + def __init__(self, isize_1=30, isize_2=48, nc=2, ndf=64, size_embed=64): super(CriticEMB, self).__init__() self.ndf = ndf self.isize_1 = isize_1 @@ -47,9 +47,9 @@ class CriticEMB(nn.Module): self.econd_lin = torch.nn.Linear(1, size_embed) # label embedding - self.fc1 = torch.nn.Linear(size_embed*3, size_embed) # 3 components after cat - self.fc2 = torch.nn.Linear(size_embed, size_embed - 32) - self.fc3 = torch.nn.Linear(size_embed - 32, 1) + self.fc1 = torch.nn.Linear(size_embed*3, size_embed*2) # 3 components after cat + self.fc2 = torch.nn.Linear(size_embed*2, size_embed) + self.fc3 = torch.nn.Linear(size_embed, 1) def forward(self, img_ECAL, img_HCAL, E_true): diff --git a/models/generatorFull.py b/models/generatorFull.py index 6178d92d0fb8840e910fb60e59845444be01a569..d8c2d7d6a58da0ec2c8ded4130b528db0b97e9f2 100644 --- a/models/generatorFull.py +++ b/models/generatorFull.py @@ -9,7 +9,7 @@ class Hcal_ecalEMB(nn.Module): """ generator component of WGAN """ - def __init__(self, ngf, ndf, nz, emb_size): + def __init__(self, ngf, ndf, nz, emb_size=16): super(Hcal_ecalEMB, self).__init__() @@ -26,11 +26,11 @@ class Hcal_ecalEMB(nn.Module): self.conv_lin_ECAL = torch.nn.Linear(7*7*7*ndf, 64) - self.econd_lin = torch.nn.Linear(1, 8) # label embedding + self.econd_lin = torch.nn.Linear(1, 64) # label embedding - self.fc1 = torch.nn.Linear(72, 64) # 2 components after cat - self.fc2 = torch.nn.Linear(64, 32) - self.fc3 = torch.nn.Linear(32, emb_size) + self.fc1 = torch.nn.Linear(64*2, 128) # 2 components after cat + self.fc2 = torch.nn.Linear(128, 64) + self.fc3 = torch.nn.Linear(64, emb_size) ## HCAL component of convolutions diff --git a/wganHCAL.py b/wganHCAL.py index 683a7568ef52edbbaf2ed84cd3c0d4f906ced667..c32b557650e1ab72c5cc7b03e6f43b9dd8ff3c46 100644 --- a/wganHCAL.py +++ b/wganHCAL.py @@ -274,8 +274,8 @@ def run(args): ## HCAL Generator and critic - mCrit = CriticEMB(args.ndf, size_embed=64).to(device) - mGen = Hcal_ecalEMB(args.ngf, 32, args.nz, emb_size=32).to(device) + mCrit = CriticEMB().to(device) + mGen = Hcal_ecalEMB(args.ngf, 32, args.nz).to(device) ## ECAL GENERATOR mGenE = DCGAN_G(args.ngf, args.nz).to(device)