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, 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, 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, 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), )