def model(self, exp):
dir = os.getcwd()
if exp == 'inversion_16':
model_dir = dir + f'/Channel/generative_model/decoder_16_VAE_1_epoch30.pth'
model = Decoder()
model.load_state_dict(
torch.load(model_dir, map_location=self.device), False)
model = model.to(self.device)
model.eval()
if exp == 'inversion_16_64':
model_dir = dir + f'/Channel/generative_model/decoder_16_64_VAE_2.5_epoch50.pth'
model = Decoder(input_channels=2)
model.load_state_dict(
torch.load(model_dir, map_location=self.device), False)
model = model.to(self.device)
model.eval()
if exp == 'inversion_16_32':
model_dir = dir + f'/Channel/generative_model/decoder_16_32_VAE_1_0.7_epoch50.pth'
model = Decoder(input_channels=2)
model.load_state_dict(
torch.load(model_dir, map_location=self.device), False)
model = model.to(self.device)
model.eval()
if exp == 'inversion_16_32_64':
model_dir = dir + f'/Channel/generative_model/decoder_16_32_64_VAE_1_0.7_0.7_epoch50.pth'
model = Decoder(input_channels=3)
model.load_state_dict(
torch.load(model_dir, map_location=self.device), False)
model = model.to(self.device)
model.eval()
return model
args = parser.parse_args()
dir = os.getcwd()
directory = f'/Gaussian/experiments/experiments_64/latent256/beta_{args.beta_vae}'
exp_dir = dir + directory + "/N{}_Bts{}_Eps{}_lr{}".\
format(args.n_train, args.batch_size, args.n_epochs, args.lr)
output_dir = exp_dir + "/save_model"
if not os.path.exists(output_dir):
os.makedirs(output_dir)
device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
encoder = Encoder()
decoder = Decoder()
encoder.to(device)
decoder.to(device)
print("number of parameters: {}".format(encoder._n_parameters() +
decoder._n_parameters()))
train_hdf5_file = os.getcwd() + \
f'/Gaussian/data/training_set_64_gaussian1_25000.hdf5'
train_loader = load_data_1scale(train_hdf5_file,
args.n_train,
args.batch_size,
singlescale=True)
optimizer = torch.optim.Adam(itertools.chain(encoder.parameters(),
decoder.parameters()),
lr=args.lr,
betas=(args.beta1, args.beta2))
