Add GMixer-24 MLP model weights, trained w/ TPU + PyTorch XLA

Author: rwightman

timm/models/mlp_mixer.py

@@ -93,7 +93,9 @@ def _cfg(url='', **kwargs):
  ),
 
  gmixer_12_224=_cfg(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD),
- gmixer_24_224=_cfg(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD),
+ gmixer_24_224=_cfg(
+ url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/gmixer_24_224_raa-7daf7ae6.pth',
+ mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD),
 
  resmlp_12_224=_cfg(
  url='https://dl.fbaipublicfiles.com/deit/resmlp_12_no_dist.pth',
@@ -457,23 +459,23 @@ def mixer_b16_224_miil_in21k(pretrained=False, **kwargs):
 
 @register_model
 def gmixer_12_224(pretrained=False, **kwargs):
- """ Glu-Mixer-12 224x224 (short & fat)
+ """ Glu-Mixer-12 224x224
  Experiment by Ross Wightman, adding (Si)GLU to MLP-Mixer
  """
  model_args = dict(
- patch_size=16, num_blocks=12, embed_dim=512, mlp_ratio=(1.0, 6.0),
+ patch_size=16, num_blocks=12, embed_dim=384, mlp_ratio=(1.0, 4.0),
  mlp_layer=GluMlp, act_layer=nn.SiLU, **kwargs)
  model = _create_mixer('gmixer_12_224', pretrained=pretrained, **model_args)
  return model
 
 
 @register_model
 def gmixer_24_224(pretrained=False, **kwargs):
- """ Glu-Mixer-24 224x224 (tall & slim)
+ """ Glu-Mixer-24 224x224
  Experiment by Ross Wightman, adding (Si)GLU to MLP-Mixer
  """
  model_args = dict(
- patch_size=16, num_blocks=24, embed_dim=384, mlp_ratio=(1.0, 6.0),
+ patch_size=16, num_blocks=24, embed_dim=384, mlp_ratio=(1.0, 4.0),
  mlp_layer=GluMlp, act_layer=nn.SiLU, **kwargs)
  model = _create_mixer('gmixer_24_224', pretrained=pretrained, **model_args)
  return model