def Featurizer(input_shape, hparams):
"""Auto-select an appropriate featurizer for the given input shape."""
if len(input_shape) == 1:
return MLP(input_shape[0], 128, hparams)
elif input_shape[1:3] == (28, 28):
return MNIST_CNN(input_shape)
elif input_shape[1:3] == (32, 32):
return wide_resnet.Wide_ResNet(input_shape, 16, 2, 0.)
elif input_shape[1:3] == (224, 224) and hparams['backbone'] in [
'resnet50', 'resnet18'
]:
return ResNet(input_shape, hparams)
elif input_shape[1:3] == (224, 224) and 'ViT-' in hparams['backbone']:
return vision_transformer.ViT2(input_shape, hparams)
elif input_shape[1:3] == (224, 224) and hparams['backbone'] in [
'B_16', 'B_32', 'L_16', 'L_32'
]:
return vision_transformer.ViT(input_shape, hparams)
elif input_shape[1:3] == (224, 224) and 'dino' in hparams['backbone']:
return vision_transformer.DINO(input_shape, hparams)
elif input_shape[1:3] == (224, 224) and 'DeiT' in hparams['backbone']:
return vision_transformer.DeiT(input_shape, hparams)
elif input_shape[1:3] == (224, 224) and 'HViT' in hparams['backbone']:
return vision_transformer.HybridViT(input_shape, hparams)
elif input_shape[1:3] == (224, 224) and 'Mixer' in hparams['backbone']:
return mlp_mixer.MLPMixer(input_shape, hparams)
elif input_shape[1:3] == (224, 224) and 'BiT' in hparams['backbone']:
return big_transfer.BiT(input_shape, hparams)
else:
raise NotImplementedError
def Featurizer(input_shape, hparams):
"""Auto-select an appropriate featurizer for the given input shape."""
if input_shape == (2048,):
return MLP(2048, 128, hparams)
elif input_shape[1:3] == (28, 28):
return MNIST_CNN(input_shape)
elif input_shape == (3, 32, 32):
return wide_resnet.Wide_ResNet(16, 2, 0.)
elif input_shape == (3, 224, 224):
return ResNet50(hparams)
def Featurizer(input_shape, hparams):
"""Auto-select an appropriate featurizer for the given input shape."""
if len(input_shape) == 1:
return MLP(input_shape[0], 128, hparams)
elif input_shape[1:3] == (28, 28):
return MNIST_CNN(input_shape)
elif input_shape[1:3] == (32, 32):
return wide_resnet.Wide_ResNet(input_shape, 16, 2, 0.)
elif input_shape[1:3] == (224, 224):
return ResNet(input_shape, hparams)
else:
raise NotImplementedError