Modules
veni.function module
- class veni.function.LeakyReLu
Bases:
veni.module.Activation- forward(x, params=None)
- generate_parameters()
- class veni.function.LogSigmoid
Bases:
veni.module.Activation- forward(x, params=None)
- generate_parameters()
- class veni.function.LogSoftmax
Bases:
veni.module.Activation- forward(x, params=None)
- generate_parameters()
- class veni.function.ReLU
Bases:
veni.module.Activation- forward(x, params=None)
- generate_parameters()
- class veni.function.Sigmoid
Bases:
veni.module.Activation- forward(x, params=None)
- generate_parameters()
- class veni.function.Softmax
Bases:
veni.module.Activation- forward(x, params=None)
- generate_parameters()
- class veni.function.Softplus
Bases:
veni.module.Activation- forward(x, params=None)
- generate_parameters()
- class veni.function.Tanh
Bases:
veni.module.Activation- forward(x, params=None)
- generate_parameters()
veni.functiontools module
- veni.functiontools.CrossEntropy(y, y_hat)
CrossEntropy loss EXPECTS: tensor of the shape (N, k1, k2, …, kn) where N is the number of examples in the batch
- Parameters
y (jnp.array) – Ground truth tensor
y_hat (jnp.array) – Model predictions
- Returns
Loss for each batch
- Return type
float
- veni.functiontools.LazyCrossEntropy(y, y_hat)
CrossEntropy loss This cross entropy implementationmay suffer numerical instabilities depending on the specific problem, consider using ‘CrossEntropy’.
EXPECTS: tensor of the shape (N, k1, k2, …, kn) where N is the number of examples in the batch
- Parameters
y (jnp.array) – Ground truth tensor
y_hat (jnp.array) – Model predictions
- Returns
Loss for each batch
- Return type
float
- veni.functiontools.MSE(y, y_hat)
Mean square error loss, reduction mean
- Parameters
y (jnp.array) – Ground truth tensor
y_hat (jnp.array) – Model predictions
- Returns
Loss for each batch
- Return type
float
- veni.functiontools.leaky_relu(x)
Applies the leaky rectified linear unit function element-wise
- Parameters
x (jax.array) – input
negative_slope (float, optional) – negative slope, defaults to 0.01
- Returns
leaky rectified linear unit on x
- Return type
jax.array
- veni.functiontools.log_sigmoid(x)
Applies the logarithmic sigmoid function element-wise
- Parameters
x (jax.array) – input
- Returns
logarithmic sigmoid on x
- Return type
jax.array
- veni.functiontools.log_softmax(x)
Applies the logarithmic softmax function element-wise.
- Parameters
x (jax.array) – input
- Returns
logarithmic softmax on x
- Return type
jax.array
- veni.functiontools.relu(x)
Applies the rectified linear unit function element-wise
- Parameters
x (jax.array) – input
- Returns
rectified linear unit on x
- Return type
jax.array
- veni.functiontools.sigmoid(x)
Applies the sigmoid function element-wise
- Parameters
x (jax.array) – input
- Returns
sigmoid on x
- Return type
jax.array
- veni.functiontools.softmax(x)
Applies the softmax function element-wise.
- Parameters
x (jax.array) – input
- Returns
softmax on x
- Return type
jax.array
- veni.functiontools.softplus(x)
Applies the softplus function element-wise. For numerical stability the implementation reverts to the linear function when input*beta > threshold.
- Parameters
x (jax.array) – input
beta (int, optional) – paramter, defaults to 1
threshold (int, optional) – threshold, defaults to 20
- Raises
ValueError – beta value must be greater than zero
- Returns
softplus on x
- Return type
jax.array
- veni.functiontools.tanh(x)
Applies the tanh function element-wise
- Parameters
x (jax.array) – input
- Returns
tanh on x
- Return type
jax.array
veni.module module
- class veni.module.Activation(f)
Bases:
abc.ABC- abstract forward(x, params=None)
- abstract generate_parameters()
- class veni.module.Sampler
Bases:
abc.ABC
veni.net module
- class veni.net.AvgPool2D(kernel_size, stride=None, padding=None)
Bases:
veni.module.Module- forward(x, params=None)
Public forward method for Conv layer
- Parameters
params (jnp.array) – Parameters of the layer
x (jnp.array) – Input
- Returns
Activation
- Return type
jnp.array
- generate_parameters()
Generate parameters for current layer
- Returns
weight and bias tensors N(0,1) initialized
- Return type
jnp.array
- property input
- property key
- property output
- class veni.net.Conv2D(inChannels, outChannels, kernelSize, stride, padding, key)
Bases:
veni.module.Module- forward(x, params)
Public forward method for Conv layer
EXPECTS: x: tensor of the form NCHW (images)x(channels)x)(height)x(width) params[0]: tensor of the form OIHW (outputCh)x(inputCh)x(kernelHeight)x(kernelWidth) params[1]: bias
- Parameters
params (jnp.array) – Parameters of the layer
x (jnp.array) – Input
- Returns
Activation
- Return type
jnp.array
- generate_parameters()
Generate parameters for current layer
- Returns
weight and bias tensors N(0,1) initialized
- Return type
jnp.array
- property input
- property key
- property output
- class veni.net.Flatten
Bases:
veni.module.Module- forward(x, params=None)
returns flattened tensor
- Returns
_description_
- Return type
_type_
TODO: optimize that
- generate_parameters()
Generate parameters for current layer
- Returns
weight and bias tensors N(0,1) initialized
- Return type
jnp.array
- property input
- property key
- property output
- class veni.net.Linear(input, output, key, bias=True)
Bases:
veni.module.Module- forward(x, params)
Public forward method for Linear layer
- Parameters
params (jnp.array) – Parameters of the layer
x (jnp.array) – Input
- Returns
Activation
- Return type
jnp.array
- generate_parameters()
- property input
- property key
- property output
- class veni.net.MLP(layers, func, key)
Bases:
veni.module.Module- forward(x, params)
- generate_parameters()
- property key
- property layers
- single_forward(x, params)
- class veni.net.MaxPool2D(kernel_size, stride=None, padding=None)
Bases:
veni.module.Module- forward(x, params=None)
Public forward method for Conv layer
- Parameters
params (jnp.array) – Parameters of the layer
x (jnp.array) – Input
- Returns
Activation
- Return type
jnp.array
- generate_parameters()
Generate parameters for current layer
- Returns
weight and bias tensors N(0,1) initialized
- Return type
jnp.array
- property input
- property key
- property output
- class veni.net.Sequential(list)
Bases:
veni.module.Module- forward(x, params)
Forward method for sequential object
- Parameters
params (jnp.array) – _description_
x (jnp.array) – _description_
- Returns
activation
- Return type
jnp.array
- generate_parameters()
Generate parameters for layers in sequential
- Returns
_description_
- Return type
jnp.array
veni.optim module
- class veni.optim.Adam(params, beta1=0.9, beta2=0.999, eta=0.001)
Bases:
veni.module.Optimizer- update(params, grads)
Update method for Adam
- Parameters
params (jax.array) – paramters to optimize
grad (jax.array) – loss gradient
- Returns
optimized parameters
- Return type
jax.array
- class veni.optim.NormalLikeSampler
Bases:
veni.module.Sampler
- class veni.optim.RademacherLikeSampler
Bases:
veni.module.Sampler
- class veni.optim.SGD(params, momentum=0, dampening=0, eta=0.001)
Bases:
veni.module.Optimizer- update(params, grad)
Update method for SGD
- Parameters
params (jax.array) – paramters to optimize
grad (jax.array) – loss gradient
- Returns
optimized parameters
- Return type
jax.array
- class veni.optim.TruncatedNormalLikeSampler(lower=- 1, upper=1)
Bases:
veni.module.Sampler
- class veni.optim.UniformLikeSampler
Bases:
veni.module.Sampler
- veni.optim.grad_fwd(params, x, y, loss, dirs=1, sampler=<veni.optim.NormalLikeSampler object>)
Function to calculate the gradient in forward mode using 1 or more directions
- Parameters
params (List) – Parameters of the model
x (jnp.array) – Input of the model
y (jnp.array) – labels
loss (Callable) – loss function
dirs (int, optional) – Number of directions used to calculate the gradient, defaults to 1
sampler (Class, optional) – Sampler used to sample gradient direction for each layer, defaults to NormalLikeSampler()
- Returns
Gradient as list of all components for each layer
- Return type
List
- veni.optim.plist_reduce(vs, js)
Multiply the jacobian vector product with the tangent directions
- Parameters
vs (list(tuple(jnp.array, jnp.array))) – tangent directions
js (float) – jacobian vector product
- Returns
multiply the jacobian vector product with the tangent directions
- Return type
list(tuple(jnp.array, jnp.array))
veni.utils module
- class veni.utils.FlattenAndCast
Bases:
object
- class veni.utils.NumpyLoader(dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None)
Bases:
Generic[torch.utils.data.dataloader.T_co]- batch_size: Optional[int]
- dataset: torch.utils.data.dataset.Dataset[torch.utils.data.dataloader.T_co]
- drop_last: bool
- num_workers: int
- pin_memory: bool
- pin_memory_device: str
- prefetch_factor: int
- sampler: Union[torch.utils.data.sampler.Sampler, Iterable]
- timeout: float
- veni.utils.numpy_collate(batch)
- veni.utils.one_hot(x, k, dtype=<class 'jax.numpy.float32'>)
Create a one-hot encoding of x of size k.