NetKet
How to use init_fun and apply_fun in MCState #1478
Replies: 1 comment 1 reply
-
|
So a few things:
For that to happen, however, you should provide me some code that works to initialise the network with the dropout layer and how to evaluate it during/not during training (With what kwargs to |
Beta Was this translation helpful? Give feedback.
-
|
Thank you for the reply. So maybe a code like this? import netket as nk
import jax
import jax.nn as jnn
import flax.linen as nn
import jax.numpy as jnp
import numpy as np
L=10
hi= nk.hilbert.Spin(s=0.5,N=L)
class FFNN(nn.Module):
alpha: int
@nn.compact
def __call__(self,x,training:bool):
some_param= self.param('some_param',nn.initializers.zeros_init(),(1,))
dropout_rng= self.make_rng('dropout')
x= nn.Dense(features= self.alpha*L, param_dtype=np.complex128, kernel_init= jnn.initializers.normal(0.01), bias_init=jnn.initializers.normal(stddev=0.01), name="layer1")(x)
x= nk.nn.reim_relu(x)
x= nn.Dropout(0.5, deterministic= not training)(x)
x= nn.Dense(features= self.alpha*L, param_dtype=np.complex128, kernel_init= jnn.initializers.normal(0.01), bias_init=jnn.initializers.normal(stddev=0.01), name="layer2")(x)
x= nk.nn.reim_relu(x)
x= jnp.sum(x, axis=(-1))
return x
model= FFNN(alpha=1)
x= hi.random_state(key=jax.random.PRNGKey(1), size=2**11)
params= model.init(rngs={'params':jax.random.PRNGKey(0), 'dropout':jax.random.PRNGKey(2)}, x=x, training=False)
y= model.apply(rngs={'dropout': jax.random.PRNGKey(5)},variables=params, x=x, training=True) |
Beta Was this translation helpful? Give feedback.
-
Hello,
I am trying to add a dropout layer, for which I need to pass the rng to the apply function of the model. But I cannot figure out how to pass the init_fun and apply_fun to the
MCState. For example (even without the dropout):I have the following error:
Click me
While I do get the output for
y=model.apply(variables=params,x=x), what would be the correct way to pass these functions toMCState?Thank you.
Beta Was this translation helpful? Give feedback.
All reactions