def __init__(self,
shape: List[int],
initializer: TensorInitArgType,
device: Optional[str] = None):
super().__init__(shape)
device = device or current_device()
add_parameter(self, 'value',
variable(shape, initializer=initializer, device=device))
def __init__(
self,
num_features,
event_ndims: int = 1,
axis: int = -1,
w_init: TensorInitArgType = init.normal,
b_init: TensorInitArgType = init.zeros,
u_init: TensorInitArgType = init.normal,
):
super().__init__(axis=axis,
event_ndims=event_ndims,
explicitly_invertible=False)
add_parameter(
self,
'w',
value=variable([1, num_features], initializer=w_init),
)
add_parameter(self, 'b', value=variable([1], initializer=b_init))
add_parameter(self,
'u',
value=variable([1, num_features], initializer=u_init))
self.num_features = num_features
self.u_hat = None
def __init__(self,
shape: List[int],
initializer: TensorInitArgType,
norm_axis: int = 1,
device: Optional[str] = None,
epsilon: float = EPSILON):
super().__init__(shape)
self.norm_axis = norm_axis
device = device or current_device()
self.epsilon = epsilon
weight = variable(shape, initializer=initializer, device=device)
with torch.no_grad():
v, _ = weight_norm_decompose(weight, norm_axis, epsilon)
add_parameter(self, 'v', v)
def __init__(self,
num_features: int,
strict: bool = False,
weight_init: TensorInitArgType = initializer.kaming_uniform,
dtype: str = 'float32',
epsilon: float = 1e-5):
"""
Construct a new linear transformation flow.
Args:
num_features: The number of features to be transformed.
The invertible transformation matrix will have the shape
``[num_features, num_features]``.
strict: Whether or not to use the strict invertible matrix?
Defaults to :obj:`False`. See :class:`LooseInvertibleMatrix`
and :class:`StrictInvertibleMatrix`.
weight_init: The weight initializer for the seed matrix.
dtype: The dtype of the invertible matrix.
epsilon: The infinitesimal constant to avoid having numerical issues.
"""
spatial_ndims = self._get_spatial_ndims()
super().__init__(
axis=-(spatial_ndims + 1),
event_ndims=(spatial_ndims + 1),
explicitly_invertible=True,
)
self.num_features = int(num_features)
self.strict = bool(strict)
self.epsilon = float(epsilon)
# Using the backend random generator instead of numpy generator
# will allow the backend random seed to have effect on the initialization
# step of the invertible matrix.
seed_matrix = variable(
shape=[num_features, num_features],
dtype=dtype,
initializer=weight_init,
requires_grad=False,
)
if strict:
self.invertible_matrix = StrictInvertibleMatrix(
to_numpy(seed_matrix), dtype=dtype, epsilon=epsilon)
else:
self.invertible_matrix = LooseInvertibleMatrix(
to_numpy(seed_matrix), dtype=dtype)
def __init__(self,
num_features: int,
axis: int = -1,
event_ndims: int = 1,
scale: Union[str, ActNormScaleType] = 'exp',
initialized: bool = False,
epsilon: float = 1e-5,
dtype: str = 'float32'):
"""
Construct a new :class:`ActNorm` instance.
Args:
num_features: The size of the feature axis.
scale: One of {"exp", "linear"}.
If "exp", ``y = (x + bias) * tf.exp(log_scale)``.
If "linear", ``y = (x + bias) * scale``.
Defaults to "exp".
axis: The axis to apply ActNorm.
Dimensions not in `axis` will be averaged out when computing
the mean of activations. Default `-1`, the last dimension.
All items of the `axis` should be covered by `event_ndims`.
event_ndims: Number of value dimensions in both `x` and `y`.
`x.ndims - event_ndims == log_det.ndims` and
`y.ndims - event_ndims == log_det.ndims`.
initialized: Whether or not the variables have been
initialized? Defaults to :obj:`False`, where the first input
`x` in the forward pass will be used to initialize the variables.
epsilon: The infinitesimal constant to avoid dividing by zero or
taking logarithm of zero.
dtype: Dtype of the parameters.
"""
# validate the arguments
scale_type = ActNormScaleType(scale)
epsilon = float(epsilon)
if scale_type == ActNormScaleType.EXP:
scale = ExpScale()
pre_scale_init = partial(init.fill, fill_value=0.)
elif scale_type == ActNormScaleType.LINEAR:
scale = LinearScale(epsilon=epsilon)
pre_scale_init = partial(init.fill, fill_value=1.)
else: # pragma: no cover
raise ValueError(f'Unsupported `scale_type`: {scale_type}')
# construct the layer
super().__init__(axis=axis,
event_ndims=event_ndims,
explicitly_invertible=True)
self.num_features = num_features
self.scale = scale
self.scale_type = scale_type.value
self.epsilon = epsilon
self.initialized = initialized
add_parameter(
self,
'pre_scale',
variable([num_features], dtype=dtype, initializer=pre_scale_init),
)
add_parameter(
self,
'bias',
variable([num_features], dtype=dtype, initializer=init.zeros),
)