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)
