Automatic EfficientNet Weight Conversion

[DavidLandup0]

Basic usage:

dummy_input_tf = tf.ones([1, 224, 224, 3])
dummy_input_torch = torch.ones(1, 3, 224, 224)

tf_model = deepvision.models.EfficientNetV2B0(include_top=False,
                                          pooling='avg',
                                          input_shape=(224, 224, 3),
                                          backend='tensorflow')

tf_model.save('effnet.h5')

from deepvision.models.classification.efficientnet import efficientnet_weight_mapper
pt_model = efficientnet_weight_mapper.load_tf_to_pt(filepath='effnet.h5', dummy_input=dummy_input_tf)

print(tf_model(dummy_input_tf)['output'].numpy())
print(pt_model(dummy_input_torch).detach().cpu().numpy())
# True
np.allclose(tf_model(dummy_input_tf)['output'].numpy(), pt_model(dummy_input_torch).detach().cpu().numpy())

pt_model = deepvision.models.EfficientNetV2B0(include_top=False,
                                          pooling='avg',
                                          input_shape=(3, 224, 224),
                                          backend='pytorch')
torch.save(pt_model.state_dict(), 'effnet.pt')

from deepvision.models.classification.efficientnet import efficientnet_weight_mapper

kwargs = {'include_top': False, 'pooling':'avg', 'input_shape':(3, 224, 224)}
tf_model = efficientnet_weight_mapper.load_pt_to_tf(filepath='effnet.pt',
                                architecture='EfficientNetV2B0',
                                kwargs=kwargs,
                                dummy_input=dummy_input_torch)


pt_model.eval()
print(pt_model(dummy_input_torch).detach().cpu().numpy())
print(tf_model(dummy_input_tf)['output'].numpy())
# True
np.allclose(tf_model(dummy_input_tf)['output'].numpy(), pt_model(dummy_input_torch).detach().cpu().numpy())