def __init__(self,
output_dim,
noise_dim=32,
input_tensor_spec=None,
hidden_layers=(256, ),
net: Network = None,
net_moving_average_rate=None,
entropy_regularization=0.,
mi_weight=None,
mi_estimator_cls=MIEstimator,
par_vi="gfsf",
optimizer=None,
name="Generator"):
r"""Create a Generator.
Args:
output_dim (int): dimension of output
noise_dim (int): dimension of noise
input_tensor_spec (nested TensorSpec): spec of inputs. If there is
no inputs, this should be None.
hidden_layers (tuple): size of hidden layers.
net (Network): network for generating outputs from [noise, inputs]
or noise (if inputs is None). If None, a default one with
hidden_layers will be created
net_moving_average_rate (float): If provided, use a moving average
version of net to do prediction. This has been shown to be
effective for GAN training (arXiv:1907.02544, arXiv:1812.04948).
entropy_regularization (float): weight of entropy regularization
mi_estimator_cls (type): the class of mutual information estimator
for maximizing the mutual information between [noise, inputs]
and [outputs, inputs].
par_vi (string): ParVI methods, options are
[``svgd``, ``svgd2``, ``svgd3``, ``gfsf``],
* svgd: empirical expectation of SVGD is evaluated by a single
resampled particle. The main benefit of this choice is it
supports conditional case, while all other options do not.
* svgd2: empirical expectation of SVGD is evaluated by splitting
half of the sampled batch. It is a trade-off between
computational efficiency and convergence speed.
* svgd3: empirical expectation of SVGD is evaluated by
resampled particles of the same batch size. It has better
convergence but involves resampling, so less efficient
computaionally comparing with svgd2.
* gfsf: wasserstein gradient flow with smoothed functions. It
involves a kernel matrix inversion, so computationally most
expensive, but in some case the convergence seems faster
than svgd approaches.
optimizer (torch.optim.Optimizer): (optional) optimizer for training
name (str): name of this generator
"""
super().__init__(train_state_spec=(), optimizer=optimizer, name=name)
self._noise_dim = noise_dim
self._entropy_regularization = entropy_regularization
self._par_vi = par_vi
if entropy_regularization == 0:
self._grad_func = self._ml_grad
else:
if par_vi == 'gfsf':
self._grad_func = self._gfsf_grad
elif par_vi == 'svgd':
self._grad_func = self._svgd_grad
elif par_vi == 'svgd2':
self._grad_func = self._svgd_grad2
elif par_vi == 'svgd3':
self._grad_func = self._svgd_grad3
else:
raise ValueError("Unsupported par_vi method: %s" % par_vi)
self._kernel_width_averager = AdaptiveAverager(
tensor_spec=TensorSpec(shape=()))
noise_spec = TensorSpec(shape=(noise_dim, ))
if net is None:
net_input_spec = noise_spec
if input_tensor_spec is not None:
net_input_spec = [net_input_spec, input_tensor_spec]
net = EncodingNetwork(input_tensor_spec=net_input_spec,
fc_layer_params=hidden_layers,
last_layer_size=output_dim,
last_activation=math_ops.identity,
name="Generator")
self._mi_estimator = None
self._input_tensor_spec = input_tensor_spec
if mi_weight is not None:
x_spec = noise_spec
y_spec = TensorSpec((output_dim, ))
if input_tensor_spec is not None:
x_spec = [x_spec, input_tensor_spec]
self._mi_estimator = mi_estimator_cls(x_spec,
y_spec,
sampler='shift')
self._mi_weight = mi_weight
self._net = net
self._predict_net = None
self._net_moving_average_rate = net_moving_average_rate
if net_moving_average_rate:
self._predict_net = net.copy(name="Genrator_average")
self._predict_net_updater = common.get_target_updater(
self._net, self._predict_net, tau=net_moving_average_rate)
def _create_averager(self):
"""Create an adaptive averager."""
return AdaptiveAverager(
tensor_spec=self._tensor_spec, speed=self._speed)
def __init__(self,
output_dim,
noise_dim=32,
input_tensor_spec=None,
hidden_layers=(256, ),
net: Network = None,
net_moving_average_rate=None,
entropy_regularization=0.,
kernel_sharpness=2.,
mi_weight=None,
mi_estimator_cls=MIEstimator,
optimizer: tf.optimizers.Optimizer = None,
name="Generator"):
"""Create a Generator.
Args:
output_dim (int): dimension of output
noise_dim (int): dimension of noise
input_tensor_spec (nested TensorSpec): spec of inputs. If there is
no inputs, this should be None.
hidden_layers (tuple): size of hidden layers.
net (Network): network for generating outputs from [noise, inputs]
or noise (if inputs is None). If None, a default one with
hidden_layers will be created
net_moving_average_rate (float): If provided, use a moving average
version of net to do prediction. This has been shown to be
effective for GAN training (arXiv:1907.02544, arXiv:1812.04948).
entropy_regularization (float): weight of entropy regularization
kernel_sharpness (float): Used only for entropy_regularization > 0.
We calcualte the kernel in SVGD as:
exp(-kernel_sharpness * reduce_mean((x-y)^2/width)),
where width is the elementwise moving average of (x-y)^2
mi_estimator_cls (type): the class of mutual information estimator
for maximizing the mutual information between [noise, inputs]
and [outputs, inputs].
optimizer (tf.optimizers.Optimizer): optimizer (optional)
name (str): name of this generator
"""
super().__init__(train_state_spec=(), optimizer=optimizer, name=name)
self._noise_dim = noise_dim
self._entropy_regularization = entropy_regularization
if entropy_regularization == 0:
self._grad_func = self._ml_grad
else:
self._grad_func = self._stein_grad
self._kernel_width_averager = AdaptiveAverager(
tensor_spec=tf.TensorSpec(shape=(output_dim, )))
self._kernel_sharpness = kernel_sharpness
noise_spec = tf.TensorSpec(shape=[noise_dim])
if net is None:
net_input_spec = noise_spec
if input_tensor_spec is not None:
net_input_spec = [net_input_spec, input_tensor_spec]
net = EncodingNetwork(
name="Generator",
input_tensor_spec=net_input_spec,
fc_layer_params=hidden_layers,
last_layer_size=output_dim)
self._mi_estimator = None
self._input_tensor_spec = input_tensor_spec
if mi_weight is not None:
x_spec = noise_spec
y_spec = tf.TensorSpec((output_dim, ))
if input_tensor_spec is not None:
x_spec = [x_spec, input_tensor_spec]
self._mi_estimator = mi_estimator_cls(
x_spec, y_spec, sampler='shift')
self._mi_weight = mi_weight
self._net = net
self._predict_net = None
self._net_moving_average_rate = net_moving_average_rate
if net_moving_average_rate:
self._predict_net = net.copy(name="Genrator_average")
tfa_common.soft_variables_update(
self._net.variables, self._predict_net.variables, tau=1.0)