def forward(self, x): output = {} transform_output = self.transformer(x) output["transform"] = transform_output["transform"] output["transform_params"] = transform_output["params"] grid = coordinates.identity_grid([self.insize, self.insize], device=x.device) grid = grid.expand(x.shape[0], *grid.shape) transformed_grid = output["transform"][-1](grid) x = x - 0.222 x = x / 0.156 view = T.broadcasting_grid_sample(x, transformed_grid) split = self.net(view) split = self.relu(split) reshape = split.view(-1, self.linear_size) z_loc = self.loc(reshape) z_scale = torch.exp(self.scale(reshape)) output["z_mu"] = z_loc output["z_std"] = z_scale return output, split
def model(self, data, transforms=None): output_size = self.encoder.insize decoder = pyro.module("decoder", self.decoder) # decoder takes z and std in the transformed coordinate frame # and the theta # and outputs an upright image with pyro.plate(data.shape[0]): # prior for z z = pyro.sample( "z", D.Normal( torch.zeros(decoder.z_dim, device=data.device), torch.ones(decoder.z_dim, device=data.device), ).to_event(1), ) # given a z, the decoder produces an "image" # this image must be transformed from the self consistent basis # to real world basis # first, z and std for the self consistent basis is outputted # then it is transfomed view = decoder(z) pyro.deterministic("canonical_view", view) # pyro.deterministic # is like pyro.sample but it is deterministic...? # all of this is completely independent of the input # maybe this is the "prior for the transformation" # and hence it looks completely independent of the input # but when the model is run again, these variables are replayed # with the theta generated by the guide # makes sense # so the model replays with theta and mu and sigma generated by # the guide, # taking theta and mu sigma and applying the inverse transform # to get the output image. grid = coordinates.identity_grid( [output_size, output_size], device=data.device) grid = grid.expand(data.shape[0], *grid.shape) transforms = T.TransformSequence(T.Translation(), T.Rotation()) transform = random_pose_transform(transforms) transform_grid = transform(grid) # output from decoder is transormed in do a different coordinate system transformed_view = T.broadcasting_grid_sample(view, transform_grid) # view from decoder outputs an image pyro.sample( "pixels", D.Bernoulli(transformed_view).to_event(3), obs=data)