How to change my jax code after I add a BatchNorm Layer?

[marcuswang6]

Due to parameter 'is_training' of BatchNorm call() , I have to change my jax code
How to deal with Jax.value_and_loss(params, x, y), model.apply(params, rng, x) after adding a BatchNorm Layer?
Where should I add this parameter 'is_training'?
Thanks in advance!

[YouJiacheng]

assume you use flax

class Model(nn.Module):
    is_training: bool
    @nn.compact
    def __call__(x):
        return SomeModuleDependsOnWhetherTraining()(x, is_training)

train_model_forward = Model(True).apply
eval_model_forward = Model(False).apply

[marcuswang6]

assume you use flax

class Model(nn.Module):
    is_training: bool
    @nn.compact
    def __call__(x):
        return SomeModuleDependsOnWhetherTraining()(x, is_training)

train_model_forward = Model(True).apply
eval_model_forward = Model(False).apply

Hi,
Thanks for your reply, JiaCheng!
I try to use flax to build a ResCNN, but there's no Conv1D block like haiku.
Now I'm using haiku, so there's no flax.nn module.
Is there any other solution for haiku user, or is there any example about Conv1D in flax?

[YouJiacheng]

haiku is similar

class MyModel(hk.Module):
    def __init__(self):
        self.bn = hk.BatchNorm(True, True, 0.99)
        self.linear = hk.Linear(16)
    def __call__(self, x, is_training):
        return self.bn(self.linear(x), is_training)

train_forward = hk.transform(lambda x: MyModel()(x, True))
eval_forward = hk.transform(lambda x: MyModel()(x, False))
params = train_forward.init(rng_key, x=x)
train_forward.apply(params, x=x, rng=rng_key)
eval_forward.apply(params, x=x, rng=rng_key)

[YouJiacheng]

You can also refer to ResNet(with BN) implementation in haiku.nets
https://github.com/deepmind/dm-haiku/blob/8e9e706c97218dcf4335400a1af04d60484d4ccd/haiku/_src/nets/resnet.py#L373-L392

[marcuswang6]

Thanks a lot~ I will check it. ^_^

[nalzok]

I think the BN layer is stateful. Don't you need haiku.transform_with_state?