ResNet-18 Initialization

[ZedongPeng]

Hi,

I am new to Flax and am trying to learn it. Regarding the ResNet example, it seems that the initial parameters, when not using pretraining, do not match those provided in torchvision.models.

from torchvision.models import resnet18
nnet = resnet18(pretrained=False)

I was wondering if using flax.linen.initializers.he_normal could help address this.

Since there are no official ResNet models in Flax, do you have any suggestions on how to replicate the ResNet-18 model from torchvision.models in JAX?

Thank you!