use_gpu = False gan_model = torch.hub.load('facebookresearch/pytorch_GAN_zoo:hub', 'DCGAN', pretrained=True, useGPU=use_gpu) from brancher.variables import ProbabilisticModel from brancher.standard_variables import NormalVariable, DeterministicVariable import brancher.functions as BF #import brancher.config as cfg #cfg.set_device('gpu') #print(cfg.device) Gmodel = gan_model.netG.cpu() decoder = BF.BrancherFunction(lambda x: Gmodel(x)) n_itr = 2000 #2000 image_size = 64 N_rep = 5 loss_list1 = [] loss_list2 = [] loss_list3 = [] imagesGT = [] imagesNoise = [] images1 = [] images2 = [] images3 = [] for rep in range(N_rep): h_size = 120 W1 = DeterministicVariable(np.random.normal(0., 0.2, (h_size, h_size)),
super(DecoderArchitecture, self).__init__() self.l1 = nn.Linear(latent_size, hidden_size) self.relu = nn.ReLU() self.l2 = nn.Linear(hidden_size, image_size) # Latent mean output self.l3 = nn.Linear(hidden_size, image_size) # Latent log sd output self.softplus = nn.Softplus() def __call__(self, x): h = self.relu(self.l1(x)) output_mean = self.l2(h) output_log_sd = self.l3(h) return {"mean": output_mean, "sd": self.softplus(output_log_sd) + 0.01} # Initialize encoder and decoders encoder = BF.BrancherFunction(EncoderArchitecture(image_size=image_size, latent_size=latent_size)) decoder = BF.BrancherFunction(DecoderArchitecture(latent_size=latent_size, image_size=image_size)) # Generative model z = NormalVariable(np.zeros((latent_size,)), np.ones((latent_size,)), name="z") decoder_output = decoder(z) x = NormalVariable(decoder_output["mean"], decoder_output["sd"], name="x") model = ProbabilisticModel([x, z]) # Amortized variational distribution Qx = EmpiricalVariable(dataset, batch_size=50, name="x", is_observed=True) encoder_output = encoder(Qx) Qz = NormalVariable(encoder_output["mean"], encoder_output["sd"], name="z") model.set_posterior_model(ProbabilisticModel([Qx, Qz])) # Joint-contrastive inference
return output_logits N_repetitions = 5 #5 num_itr = 3000 N_ELBO = 10 N_ELBO_ITR = 20 b_size = 200 loss_list1 = [] loss_list2 = [] loss_list3 = [] for rep in range(N_repetitions): # # Initialize encoder and decoders encoder1 = BF.BrancherFunction( EncoderArchitecture1(image_size=image_size, latent_size3=latent_size3)) encoder2 = BF.BrancherFunction( EncoderArchitecture2(latent_size2=latent_size2, latent_size3=latent_size3)) encoder3 = BF.BrancherFunction( EncoderArchitecture3(latent_size1=latent_size1, latent_size2=latent_size2)) decoder1 = BF.BrancherFunction( DecoderArchitecture1(latent_size1=latent_size1, latent_size2=latent_size2)) decoder2 = BF.BrancherFunction( DecoderArchitecture2(latent_size2=latent_size2, latent_size3=latent_size3)) decoder3 = BF.BrancherFunction( DecoderArchitecture3(latent_size3=latent_size3, image_size=image_size))
output_mean = self.l1(x) return {"mean": output_mean} N_repetitions = 15 num_itr = 3000 N_ELBO = 10 N_ELBO_ITR = 20 loss_list1 = [] loss_list2 = [] loss_list3 = [] for rep in range(N_repetitions): print("Repetition: {}".format(rep)) # Initialize encoder and decoders encoder1 = BF.BrancherFunction(EncoderArchitecture1(image_size=image_size, latent_size3=latent_size3)) encoder2 = BF.BrancherFunction(EncoderArchitecture2(latent_size2=latent_size2, latent_size3=latent_size3)) encoder3 = BF.BrancherFunction(EncoderArchitecture3(latent_size1=latent_size1, latent_size2=latent_size2)) decoder1 = BF.BrancherFunction(DecoderArchitecture1(latent_size1=latent_size1, latent_size2=latent_size2)) decoder2 = BF.BrancherFunction(DecoderArchitecture2(latent_size2=latent_size2, latent_size3=latent_size3)) decoder3 = BF.BrancherFunction(DecoderArchitecture3(latent_size3=latent_size3, image_size=image_size)) # Generative model z1sd = 1.5 #1 z2sd = 0.25 #0.25 z3sd = 0.15 z1 = NormalVariable(np.zeros((latent_size1,)), z1sd*np.ones((latent_size1,)), name="z1") decoder_output1 = DeterministicVariable(decoder1(z1), name="decoder_output1") z2 = NormalVariable(BF.relu(decoder_output1["mean"]), z2sd*np.ones((latent_size2,)), name="z2") decoder_output2 = DeterministicVariable(decoder2(z2), name="decoder_output2")
self.fc1 = nn.Linear(7, n_hidden) # 6*6 from image dimension self.fc2 = nn.Linear(n_hidden, 2) self.fc3 = nn.Linear(n_hidden, 2) def forward(self, r, w0, w1, mx, my, sx, sy): x = torch.cat((r, w0, w1, mx, my, sx, s), 1) x = F.relu(self.fc1(x)) x = self.fc2(x) x = F.softplus(self.fc3(x)) return x, y net = InferenceNet() # Variational bayesian update brancher_net = BF.BrancherFunction(net) def bayesian_update(r, w0, w1, mx, my, sx, sy): out = brancher_net(r, w0, w1, mx, my, sx, sy) return [out[0], out[1], out[2], out[3]] # Model T = 20 sigma = 0.1 mx = [DeterministicVariable(0., "mx_0")] my = [DeterministicVariable(0., "my_0")] sx = [DeterministicVariable(0., "sx_0")] sy = [DeterministicVariable(0., "sy_0")]
#output_log_sd = self.l3(h0) return {"mean": output_mean} class DecoderArchitecture2(nn.Module): def __init__(self, latent_size2, image_size): super(DecoderArchitecture2, self).__init__() self.l1 = nn.Linear(latent_size2, image_size) def __call__(self, x): output_mean = self.l1(x) return {"mean": output_mean} # Initialize encoder and decoders encoder1 = BF.BrancherFunction( EncoderArchitecture1(image_size=image_size, latent_size2=latent_size2)) encoder2 = BF.BrancherFunction( EncoderArchitecture2(latent_size1=latent_size1, latent_size2=latent_size2)) decoder1 = BF.BrancherFunction( DecoderArchitecture1(latent_size1=latent_size1, latent_size2=latent_size2)) decoder2 = BF.BrancherFunction( DecoderArchitecture2(latent_size2=latent_size2, image_size=image_size)) # Generative model z1sd = 1. z2sd = 0.5 #0.01 z1 = NormalVariable(np.zeros((latent_size1, )), z1sd * np.ones((latent_size1, )), name="z1") decoder_output1 = DeterministicVariable(decoder1(z1), name="decoder_output1") z2 = NormalVariable(BF.relu(decoder_output1["mean"]),
padding=1) self.f1 = torch.nn.ReLU() self.l2 = torch.nn.Linear(in_features=image_size**2 * out_channels, out_features=10) def __call__(self, x): h = self.f1(self.l1(x)) h_shape = h.shape h = h.view((h_shape[0], np.prod(h_shape[1:]))) logits = self.l2(h) return logits network = PytorchNetwork() brancher_network = BF.BrancherFunction(network) # Data sampling model minibatch_size = 4 minibatch_indices = RandomIndices(dataset_size=dataset_size, batch_size=minibatch_size, name="indices", is_observed=True) x = Empirical(input_variable, indices=minibatch_indices, name="x", is_observed=True) labels = Empirical(output_labels, indices=minibatch_indices, name="labels", is_observed=True)