Commit c4880a8c authored by MattiaPujatti's avatar MattiaPujatti
Browse files

updated classifier wrapper

parent 2de9c3c8
This diff is collapsed.
This diff is collapsed.
......@@ -21,89 +21,3 @@ def initialize_cmplx_haiku_model(model, init_shape, rng_seed=42, **model_kwargs)
return network, net_params, net_state
def evaluate_model(network,
) -> Any:
"""Calling this function one can compute the `metric` estimate over the input
dataset, provided as a Pytorch dataloader.
dataloader: pytorch DataLoader
Dataloader containing all the test samples.
log_metric = []
for batch in tqdm(dataloader, desc='Looping over the dataset.', unit='batches', leave=False):
x_batch, y_batch = batch[0].numpy(), batch[1].numpy()
batch_metric, _ = metric( model, net_params, x_batch, y_batch, next(self.rng_seq), net_state, is_training=False )
log_metric.append( batch_metric )
return np.mean(log_metric)
@partial(jit, static_argnums=(0,8,))
def update(
params: hk.Params,
x: jnp.array,
y: jnp.array,
metric: Callable,
opt_state: Any,
opt_update: Callable,
get_params: Callable,
rng_key: Optional[jax.random.PRNGKey] =None,
net_state: Optional[hk.State] = None,
is_training: bool = False,
step: int = 0,
) -> Tuple[hk.Params, Any, float, hk.State]:
"""Given a minibatch of samples, it compute the loss and the parameters updates of the network.
Then, since jax.grad calculate the complex gradient df/dz and not the conjugate (as needed by
the complex backpropagation), an additional step performs this operation, before applying the
updates just computed.
params: hk.Module.parameters()
Parameters of the network.
opt_state: jax pytree
Object representing the actual optimizer state.
x: array
Array of samples to give in input to the network.
y: array
Array of one-hot encoded labels.
rng_key: jax.random.PRNGKey, optional (default is None)
PRNGKey necessary to the 'apply' method of the transformed network.
net_state: hk.State, optional (default is None)
Internal state of the network. Set 'None' if the network has no internal trainable state.
is_training: bool, optional (default is False)
Flags that alert the network if it is called in training or evaluation mode. Useful in presence
of dropout or batchnormalization layers.
step: int, optional (default is 0)
Index of the update step.
new_params: hk.Module.parameters
New estimates of network's parameters.
opt_state: jax pytree
Optimizer state after the update.
loss: float
Loss estimate for the given minibatch.
Internal state of the network.
(loss, net_state), grads = value_and_grad(self.__crossentropy_loss, has_aux=True)(params, x, y, rng_key, net_state, is_training)
grads = jax.tree_multimap(jnp.conjugate, grads)
#print(jax.tree_multimap(jnp.mean, grads))
opt_state = self.opt_update(step, grads, opt_state)
return self.get_params(opt_state), opt_state, loss, net_state
......@@ -59,9 +59,9 @@ def categorical_accuracy(network,
@partial(jit, static_argnums=(0,6,))
def crossentropy_loss(network,
params: hk.Params,
@partial(jit, static_argnums=(1,6,))
def crossentropy_loss(params: hk.Params,
inputs: jnp.array,
targets: jnp.array,
rng_key: Optional[jax.random.PRNGKey] =None,
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment