def type(self): in_tys = [ types.tensor_type(torch.float, in_dim) for in_dim in self._in_dims ] in_space = functools.reduce(lambda t, u: t @ u, in_tys, Ty()) out_tys = [ types.tensor_type(torch.float, out_dim) for out_dim in self._out_dims ] out_space = functools.reduce(lambda t, u: t @ u, out_tys, Ty()) return in_space >> out_space
def __init__(self, data_dim=28 * 28, hidden_dim=64, guide_hidden_dim=256): self._data_dim = data_dim # Build up a bunch of torch.Sizes for the powers of two between # hidden_dim and data_dim. dims = list(util.powers_of(2, hidden_dim, data_dim // 4)) + [data_dim] dims.sort() gaussian_likelihood = lambda dim: DiagonalGaussian( dim, latent_name='X^{%d}' % dim, likelihood=True) generators = [] for lower, higher in zip(dims, dims[1:]): # Construct the VLAE decoder and encoder if higher == self._data_dim: decoder = LadderDecoder(lower, higher, noise_dim=2, conv=True, out_dist=gaussian_likelihood) else: decoder = LadderDecoder(lower, higher, noise_dim=2, conv=False, out_dist=None) data = {'effect': decoder.effect} generator = cart_closed.Box(decoder.name, decoder.type.left, decoder.type.right, decoder, data=data) generators.append(generator) # For each dimensionality, construct a prior/posterior ladder pair for dim in set(dims) - {data_dim}: space = types.tensor_type(torch.float, dim) prior = LadderPrior(dim, None) generator = cart_closed.Box(prior.name, Ty(), space, prior, data={'effect': prior.effect}) generators.append(generator) super().__init__(generators, [], data_dim, guide_hidden_dim, list(set(dims) - {data_dim}))
def __init__(self, generators, global_elements=[], data_space=(784, ), guide_hidden_dim=256, no_prior_dims=[]): super().__init__() if isinstance(data_space, int): data_space = (data_space, ) self._data_space = data_space self._data_dim = math.prod(data_space) if len(self._data_space) == 1: self._observation_name = '$X^{%d}$' % self._data_dim else: self._observation_name = '$X^{%s}$' % str(self._data_space) obs = set() for generator in generators: ty = generator.dom >> generator.cod obs = obs | unification.base_elements(ty) for element in global_elements: ty = element.dom >> element.cod obs = obs - unification.base_elements(ty) no_prior_dims = no_prior_dims + [self._data_dim] for ob in obs: dim = types.type_size(str(ob)) if dim in no_prior_dims: continue space = types.tensor_type(torch.float, dim) prior = StandardNormal(dim) name = '$p(%s)$' % prior.effects effect = {'effect': prior.effect, 'dagger_effect': []} global_element = cart_closed.Box(name, Ty(), space, prior, data=effect) global_elements.append(global_element) self._category = freecat.FreeCategory(generators, global_elements) self.guide_temperatures = nn.Sequential( nn.Linear(self._data_dim, guide_hidden_dim), nn.LayerNorm(guide_hidden_dim), nn.PReLU(), nn.Linear(guide_hidden_dim, 1 * 2), nn.Softplus(), ) self.guide_arrow_weights = nn.Sequential( nn.Linear(self._data_dim, guide_hidden_dim), nn.LayerNorm(guide_hidden_dim), nn.PReLU(), nn.Linear(guide_hidden_dim, self._category.arrow_weight_loc.shape[0] * 2), ) self._random_variable_names = collections.defaultdict(int) self.encoders = nn.ModuleDict() self.encoder_functor = wiring.Functor( lambda ty: util.double_latent(ty, self.data_space), lambda ar: self._encoder(ar.name), ob_factory=Ty, ar_factory=cart_closed.Box) for arrow in self._category.ars: effect = arrow.data['effect'] cod_dims = util.double_latents( [types.type_size(ob.name) for ob in arrow.cod], self._data_dim) dom_dims = util.double_latents( [types.type_size(ob.name) for ob in arrow.dom], self._data_dim) self.encoders[arrow.name + '†'] = build_encoder( cod_dims, dom_dims, effect)
def __init__(self, data_dim=28 * 28, hidden_dim=64, guide_hidden_dim=256): self._data_dim = data_dim data_side = int(math.sqrt(self._data_dim)) glimpse_side = data_side // 2 glimpse_dim = glimpse_side**2 generators = [] # Build up a bunch of torch.Sizes for the powers of two between # hidden_dim and glimpse_dim. dims = list(util.powers_of(2, hidden_dim, glimpse_dim // 4)) +\ [glimpse_dim] dims.sort() for dim_a, dim_b in itertools.combinations(dims, 2): lower, higher = sorted([dim_a, dim_b]) # Construct the decoder and encoder if higher == glimpse_dim: decoder = DensityDecoder(lower, higher, DiagonalGaussian, convolve=True) else: decoder = DensityDecoder(lower, higher, DiagonalGaussian) data = {'effect': decoder.effect} generator = cart_closed.Box(decoder.name, decoder.type.left, decoder.type.right, decoder, data=data) generators.append(generator) # Build up a bunch of torch.Sizes for the powers of two between # hidden_dim and data_dim. dims = dims + [data_dim] dims.sort() for lower, higher in zip(dims, dims[1:]): # Construct the VLAE decoder and encoder if higher == self._data_dim: decoder = LadderDecoder(lower, higher, noise_dim=2, conv=True, out_dist=DiagonalGaussian) else: decoder = LadderDecoder(lower, higher, noise_dim=2, conv=False, out_dist=DiagonalGaussian) data = {'effect': decoder.effect} generator = cart_closed.Box(decoder.name, decoder.type.left, decoder.type.right, decoder, data=data) generators.append(generator) # For each dimensionality, construct a prior/posterior ladder pair for dim in set(dims) - {glimpse_dim, data_dim}: space = types.tensor_type(torch.float, dim) prior = LadderPrior(dim, DiagonalGaussian) data = {'effect': prior.effect} generator = cart_closed.Box(prior.name, Ty(), space, prior, data=data) generators.append(generator) # Construct writer/reader pair for spatial attention writer = SpatialTransformerWriter(data_side, glimpse_side) writer_l, writer_r = writer.type.left, writer.type.right data = {'effect': writer.effect} generator = cart_closed.Box(writer.name, writer_l, writer_r, writer, data=data) generators.append(generator) # Construct the likelihood likelihood = GaussianLikelihood(data_dim, 'X^{%d}' % data_dim) data = {'effect': likelihood.effect} generator = cart_closed.Box(likelihood.name, likelihood.type.left, likelihood.type.right, likelihood, data=data) generators.append(generator) super().__init__(generators, [], data_dim, guide_hidden_dim, no_prior_dims=[glimpse_dim, data_dim])
def data_space(self): return types.tensor_type(torch.float, self._data_space)
def type(self): return Ty() >> types.tensor_type(torch.float, self._dim)
def type(self): canvas_type = types.tensor_type(torch.float, self._canvas_side**2) glimpse_type = types.tensor_type(torch.float, self._glimpse_side**2) return (canvas_type @ glimpse_type) >> canvas_type
def type(self): embedding_type = types.tensor_type(torch.float, self._hidden_dim) smiles_type = types.tensor_type(torch.float, (self._max_len, self._charset_len)) return embedding_type >> smiles_type
def type(self): glimpse_type = types.tensor_type(torch.float, self._glimpse_side**2) canvas_type = types.tensor_type(torch.float, self._canvas_side**2) return glimpse_type >> canvas_type
def type(self): dim_space = types.tensor_type(torch.float, self._dim) return (dim_space @ dim_space) >> dim_space
def type(self): input_space = types.tensor_type(torch.float, self._in_dim) noise_space = types.tensor_type(torch.float, self._noise_dim) return (input_space @ noise_space) >>\ types.tensor_type(torch.float, self._out_dim)
def __init__(self, in_dim, out_dim, dist_layer=ContinuousBernoulliModel, normalizer_layer=nn.LayerNorm, convolve=False): super().__init__() self._in_dim = in_dim self._out_dim = out_dim self._in_space = types.tensor_type(torch.float, in_dim) self._out_space = types.tensor_type(torch.float, out_dim) self._convolve = convolve self.add_module('distribution', dist_layer(out_dim)) hidden_dim = (in_dim + out_dim) // 2 self._channels = 1 if isinstance(self.distribution, DiagonalGaussian): self._channels *= 2 final_features = out_dim * self._channels if not self._convolve: self.add_module( 'neural_layers', nn.Sequential( nn.Linear(in_dim, hidden_dim), normalizer_layer(hidden_dim), nn.PReLU(), nn.Linear(hidden_dim, hidden_dim), normalizer_layer(hidden_dim), nn.PReLU(), nn.Linear(hidden_dim, final_features), )) else: if out_dim > in_dim: out_side = int(np.sqrt(self._out_dim)) conv_side = max(out_side // 4, 1) multiplier = conv_side**2 self.dense_layers = nn.Sequential( nn.Linear(self._in_dim, multiplier * 2 * out_side), normalizer_layer(multiplier * 2 * out_side), nn.PReLU(), nn.Linear(multiplier * 2 * out_side, multiplier * 2 * out_side), ) self.conv_layers = nn.Sequential( nn.ConvTranspose2d(2 * out_side, out_side, 4, 2, 1), nn.InstanceNorm2d(out_side), nn.PReLU(), nn.ConvTranspose2d(out_side, self._channels, 4, 2, 1), nn.ReflectionPad2d((out_side - conv_side * 4) // 2)) else: in_side = int(np.sqrt(self._in_dim)) multiplier = max(in_side // 4, 1)**2 self.conv_layers = nn.Sequential( nn.Conv2d(1, in_side, 4, 2, 1), nn.InstanceNorm2d(in_side), nn.PReLU(), nn.Conv2d(in_side, in_side * 2, 4, 2, 1), nn.InstanceNorm2d(in_side * 2), nn.PReLU(), ) self.dense_layers = nn.Sequential( nn.Linear(in_side * 2 * multiplier, final_features), normalizer_layer(final_features), nn.PReLU(), nn.Linear(final_features, final_features))
def type(self): return types.tensor_type(torch.float, self._dim) >>\ types.tensor_type(torch.float, self._dim)