def __init__(self, use_affine_ex: bool = True, **kwargs):
super().__init__(**kwargs)
self.use_affine_ex = use_affine_ex
# decoder parts
# Flow for modelling t Gamma
self.thickness_flow_components = ComposeTransformModule([Spline(1)])
self.thickness_flow_constraint_transforms = ComposeTransform(
[self.thickness_flow_lognorm,
ExpTransform()])
self.thickness_flow_transforms = ComposeTransform([
self.thickness_flow_components,
self.thickness_flow_constraint_transforms
])
# affine flow for s normal
intensity_net = DenseNN(1, [1],
param_dims=[1, 1],
nonlinearity=torch.nn.Identity())
self.intensity_flow_components = ConditionalAffineTransform(
context_nn=intensity_net, event_dim=0)
self.intensity_flow_constraint_transforms = ComposeTransform(
[SigmoidTransform(), self.intensity_flow_norm])
self.intensity_flow_transforms = [
self.intensity_flow_components,
self.intensity_flow_constraint_transforms
]
# build flow as s_affine_w * t * e_s + b -> depends on t though
# realnvp or so for x
self._build_image_flow()
def __init__(self, **kwargs):
super().__init__(**kwargs)
# Flow for modelling t Gamma
self.thickness_flow_components = ComposeTransformModule([Spline(1)])
self.thickness_flow_constraint_transforms = ComposeTransform(
[self.thickness_flow_lognorm,
ExpTransform()])
self.thickness_flow_transforms = ComposeTransform([
self.thickness_flow_components,
self.thickness_flow_constraint_transforms
])
# affine flow for s normal
intensity_net = DenseNN(1, [1],
param_dims=[1, 1],
nonlinearity=torch.nn.Identity())
self.intensity_flow_components = ConditionalAffineTransform(
context_nn=intensity_net, event_dim=0)
self.intensity_flow_constraint_transforms = ComposeTransform(
[SigmoidTransform(), self.intensity_flow_norm])
self.intensity_flow_transforms = [
self.intensity_flow_components,
self.intensity_flow_constraint_transforms
]
def __init__(self, use_affine_ex: bool = True, **kwargs):
super().__init__(**kwargs)
self.use_affine_ex = use_affine_ex
# decoder parts
# Flow for modelling t Gamma
self.thickness_flow_components = ComposeTransformModule([Spline(1)])
self.thickness_flow_constraint_transforms = ComposeTransform(
[self.thickness_flow_lognorm,
ExpTransform()])
self.thickness_flow_transforms = ComposeTransform([
self.thickness_flow_components,
self.thickness_flow_constraint_transforms
])
# affine flow for s normal
self.intensity_flow_components = ComposeTransformModule(
[LearnedAffineTransform(), Spline(1)])
self.intensity_flow_constraint_transforms = ComposeTransform(
[SigmoidTransform(), self.intensity_flow_norm])
self.intensity_flow_transforms = [
self.intensity_flow_components,
self.intensity_flow_constraint_transforms
]
# realnvp or so for x
self._build_image_flow()
def model(n_samples=None, scale=2.):
with pyro.plate('observations', n_samples):
thickness = pyro.sample('thickness', Gamma(10., 5.))
loc = (thickness - 2.5) * 2
transforms = ComposeTransform([SigmoidTransform(), AffineTransform(10, 15)])
width = pyro.sample('width', TransformedDistribution(Normal(loc, scale), transforms))
return thickness, width
def __init__(self, use_affine_ex=True, **kwargs):
super.__init__(**kwargs)
self.num_scales = 2
self.register_buffer("glasses_base_loc",
torch.zeros([
1,
], requires_grad=False))
self.register_buffer("glasses_base_scale",
torch.ones([
1,
], requires_grad=False))
self.register_buffer("glasses_flow_lognorm_loc",
torch.zeros([], requires_grad=False))
self.register_buffer("glasses_flow_lognorm_scale",
torch.ones([], requires_grad=False))
self.glasses_flow_components = ComposeTransformModule([Spline(1)])
self.glasses_flow_constraint_transforms = ComposeTransform(
[self.glasses_flow_lognorm,
SigmoidTransform()])
self.glasses_flow_transforms = ComposeTransform([
self.glasses_flow_components,
self.glasses_flow_constraint_transforms
])
glasses_base_dist = Normal(self.glasses_base_loc,
self.glasses_base_scale).to_event(1)
self.glasses_dist = TransformedDistribution(
glasses_base_dist, self.glasses_flow_transforms)
glasses_ = pyro.sample("glasses_", self.glasses_dist)
glasses = pyro.sample("glasses", dist.Bernoulli(glasses_))
glasses_context = self.glasses_flow_constraint_transforms.inv(glasses_)
self.x_transforms = self._build_image_flow()
self.register_buffer("x_base_loc",
torch.zeros([1, 64, 64], requires_grad=False))
self.register_buffer("x_base_scale",
torch.ones([1, 64, 64], requires_grad=False))
x_base_dist = Normal(self.x_base_loc, self.x_base_scale).to_event(3)
cond_x_transforms = ComposeTransform(
ConditionalTransformedDistribution(
x_base_dist,
self.x_transforms).condition(context).transforms).inv
cond_x_dist = TransformedDistribution(x_base_dist, cond_x_transforms)
x = pyro.sample("x", cond_x_dist)
return x, glasses
def _get_preprocess_transforms(self):
alpha = 0.05
num_bits = 8
if self.preprocessing == 'glow':
# Map to [-0.5,0.5]
a1 = AffineTransform(-0.5, (1. / 2 ** num_bits))
preprocess_transform = ComposeTransform([a1])
elif self.preprocessing == 'realnvp':
# Map to [0,1]
a1 = AffineTransform(0., (1. / 2 ** num_bits))
# Map into unconstrained space as done in RealNVP
a2 = AffineTransform(alpha, (1 - alpha))
s = SigmoidTransform()
preprocess_transform = ComposeTransform([a1, a2, s.inv])
else:
raise ValueError(f'{self.preprocessing} not valid.')
return preprocess_transform
def model(n_samples=None, scale=0.5, invert=False):
with pyro.plate('observations', n_samples):
thickness = 0.5 + pyro.sample('thickness', Gamma(10., 5.))
if invert:
loc = (thickness - 2) * -2
else:
loc = (thickness - 2.5) * 2
transforms = ComposeTransform(
[SigmoidTransform(), AffineTransform(64, 191)])
intensity = pyro.sample(
'intensity', TransformedDistribution(Normal(loc, scale),
transforms))
return thickness, intensity
def __init__(self, **kwargs):
super().__init__(**kwargs)
# Flow for modelling t Gamma
self.thickness_flow_components = ComposeTransformModule([Spline(1)])
self.thickness_flow_constraint_transforms = ComposeTransform(
[self.thickness_flow_lognorm,
ExpTransform()])
self.thickness_flow_transforms = ComposeTransform([
self.thickness_flow_components,
self.thickness_flow_constraint_transforms
])
# affine flow for s normal
self.intensity_flow_components = ComposeTransformModule(
[LearnedAffineTransform(), Spline(1)])
self.intensity_flow_constraint_transforms = ComposeTransform(
[SigmoidTransform(), self.intensity_flow_norm])
self.intensity_flow_transforms = [
self.intensity_flow_components,
self.intensity_flow_constraint_transforms
]
def sigmoid(self):
return self.transform(SigmoidTransform())