def __init__(self,
num_inputs,
num_hidden,
mask,
num_cond_inputs=None,
act='relu',
normalization_str='batchnorm',
layer_modifier='none'):
super(AdditiveCouplingLayer, self).__init__()
self.num_inputs = num_inputs
self.mask = mask
act_func = layers.str_to_activ_module(act)
if num_cond_inputs is not None:
total_inputs = num_inputs + num_cond_inputs
else:
total_inputs = num_inputs
# Grab the type of linear layer
layer_fn = get_layer_fn(layer_modifier)
self.net = nn.Sequential(
layers.BasicDenseBlock(total_inputs, num_hidden, layer_fn=layer_fn,
normalization_str=normalization_str, activation_str=act),
act_func(),
# nn.Linear(total_inputs, num_hidden), act_func(),
layers.BasicDenseBlock(num_hidden, num_hidden, layer_fn=layer_fn,
normalization_str=normalization_str, activation_str=act),
act_func(),
# nn.Linear(num_hidden, num_hidden), act_func(),
nn.Linear(num_hidden, num_inputs))
def __init__(self, vae, output_size, **kwargs):
super(Saccader, self).__init__()
self.vae = vae
self.output_size = output_size
self.config = kwargs['kwargs']
# build the projection to softmax from RNN state
self.latent_projector = ImageStateProjector(
output_size=self.output_size, config=self.config)
# build the inlay projector
self.numgrad = NumericalDifferentiator(self.config)
# build the spatial transformer (or local spatial transformer)
self.spatial_transformer = SpatialTransformer(self.config) \
if self.config['task'] != 'crop_dual_imagefolder' \
else LocalizedSpatialTransformer(self.vae.chans, self.config)
# build the projector from the posterior to the ST params
self.posterior_to_st = nn.Sequential(
self._get_dense('posterior_to_st_proj')(
self.vae.reparameterizer.output_size,
3,
normalization_str=self.config['dense_normalization'],
activation_fn=str_to_activ_module(self.config['activation']))
) if self.vae.reparameterizer.output_size != 3 else Identity()
def _append_inter_projection(self, reparam, input_size, output_size, name):
return nn.Sequential(
self._get_dense_net_map(name=name)(
input_size,
output_size,
nlayers=2,
activation_fn=str_to_activ_module(self.config['activation'])),
reparam)
def __init__(self,
num_inputs,
num_hidden,
num_cond_inputs=None,
s_act='tanh',
t_act='relu',
pre_exp_tanh=False):
super(MADESplit, self).__init__()
self.pre_exp_tanh = pre_exp_tanh
input_mask = get_mask(num_inputs, num_hidden, num_inputs,
mask_type='input')
hidden_mask = get_mask(num_hidden, num_hidden, num_inputs)
output_mask = get_mask(num_hidden, num_inputs, num_inputs,
mask_type='output')
act_func = layers.str_to_activ_module(s_act)
self.s_joiner = MaskedLinear(num_inputs, num_hidden, input_mask,
num_cond_inputs)
self.s_trunk = nn.Sequential(act_func(),
MaskedLinear(num_hidden, num_hidden,
hidden_mask), act_func(),
MaskedLinear(num_hidden, num_inputs,
output_mask))
act_func = layers.str_to_activ_module(t_act)
self.t_joiner = MaskedLinear(num_inputs, num_hidden, input_mask,
num_cond_inputs)
self.t_trunk = nn.Sequential(act_func(),
MaskedLinear(num_hidden, num_hidden,
hidden_mask), act_func(),
MaskedLinear(num_hidden, num_inputs,
output_mask))
def __init__(self, input_shape, **kwargs):
super(AbstractVAE, self).__init__()
self.input_shape = input_shape
self.is_color = input_shape[0] > 1
self.chans = 3 if self.is_color else 1
# grab the meta config and print for
self.config = kwargs['kwargs']
pp = pprint.PrettyPrinter(indent=4)
pp.pprint(self.config)
# grab the activation nn.Module from the string
self.activation_fn = str_to_activ_module(self.config['activation'])
# placeholder in order to sequentialize model
self.full_model = None
def __init__(self, input_shape, **kwargs):
""" Abstract base class for VAE.
:param input_shape: the input tensor shape
:returns: instantiation of object
:rtype: object
"""
super(AbstractVAE, self).__init__()
self.input_shape = input_shape
self.is_color = input_shape[0] > 1
self.chans = 3 if self.is_color else 1
self.config = kwargs['kwargs']
# keep track of ammortized posterior
self.aggregate_posterior = EMA(0.999)
# grab the activation nn.Module from the string
self.activation_fn = str_to_activ_module(self.config['activation'])
def __init__(self,
num_inputs,
num_hidden,
num_cond_inputs=None,
act='relu',
pre_exp_tanh=False):
super(MADE, self).__init__()
act_func = layers.str_to_activ_module(act)
input_mask = get_mask(
num_inputs, num_hidden, num_inputs, mask_type='input')
hidden_mask = get_mask(num_hidden, num_hidden, num_inputs)
output_mask = get_mask(
num_hidden, num_inputs * 2, num_inputs, mask_type='output')
self.joiner = MaskedLinear(num_inputs, num_hidden, input_mask,
num_cond_inputs)
self.trunk = nn.Sequential(act_func(),
MaskedLinear(num_hidden, num_hidden,
hidden_mask), act_func(),
MaskedLinear(num_hidden, num_inputs * 2,
output_mask))
def __init__(self, input_shape, n_layers=2, bidirectional=False, **kwargs):
""" Implementation of the Variational Recurrent
Neural Network (VRNN) from https://arxiv.org/abs/1506.02216
:param input_shape: the input dimension
:param n_layers: number of RNN / equivalent layers
:param bidirectional: whether the model is bidirectional or not
:returns: VRNN object
:rtype: AbstractVAE
"""
super(VRNN, self).__init__(input_shape, **kwargs)
assert self.config[
'max_time_steps'] > 0, "Need max_time_steps > 0 for VRNN."
self.activation_fn = str_to_activ_module(self.config['activation'])
self.bidirectional = bidirectional
self.n_layers = n_layers
# build the reparameterizer
self.reparameterizer = get_reparameterizer(
self.config['reparam_type'])(self.config)
# keep track of ammortized posterior
self.aggregate_posterior = nn.ModuleDict({
'encoder_logits':
EMA(self.config['aggregate_posterior_ema_decay']),
'prior_logits':
EMA(self.config['aggregate_posterior_ema_decay']),
'rnn_hidden_state_h':
EMA(self.config['aggregate_posterior_ema_decay']),
'rnn_hidden_state_c':
EMA(self.config['aggregate_posterior_ema_decay'])
})
# build the entire model
self._build_model()
def _build_model(self):
''' helper function to build convolutional or dense decoder
chans * 2 because we want to do relationships'''
crop_size = [
self.config['img_shp'][0], self.config['window_size'],
self.config['window_size']
]
# main function approximator to extract crop features
nlayer_map = {
'conv': {
32: 4,
70: 6,
100: 7
},
'resnet': {
32: None,
64: None,
70: None,
100: None
},
'dense': {
32: 3,
64: 3,
70: 3,
100: 3
}
}
bilinear_size = (self.config['window_size'],
self.config['window_size'])
conv = get_encoder(self.config, name='crop_feat')(
crop_size,
self.config['latent_size'],
activation_fn=str_to_activ_module(self.config['activation']),
num_layers=nlayer_map[self.config['encoder_layer_type']][
self.config['window_size']],
bilinear_size=bilinear_size)
# takes the state + output of conv and projects it
# the +1 is for ACT
state_output_size = self.config['latent_size'] + 1 if self.config['concat_prediction_size'] <= 0 \
else self.config['concat_prediction_size'] + 1
state_input_size = self.config['latent_size'] if self.config['disable_rnn_proj'] \
else self.config['latent_size']*2
state_projector = nn.Sequential(
self._get_dense(name='state_proj')(
state_input_size,
state_output_size,
normalization_str=self.config['dense_normalization'],
activation_fn=str_to_activ_module(self.config['activation'])))
# takes the finally aggregated vector and projects to output dims
input_size = self.config['concat_prediction_size'] * self.config['max_time_steps'] \
if self.config['concat_prediction_size'] > 0 else self.config['latent_size']
output_projector = self._get_dense(name='output_proj')(
input_size,
self.output_size,
normalization_str=self.config['dense_normalization'],
activation_fn=str_to_activ_module(self.config['activation']))
return conv, state_projector, output_projector