├── .gitignore ├── LICENSE ├── README.md ├── docker ├── pt_git.Dockerfile ├── pt_pip.Dockerfile ├── tf_git.Dockerfile ├── tf_git_ampere.Dockerfile └── tf_pip.Dockerfile ├── export_pytorch.py ├── jeffnet ├── __init__.py ├── common │ ├── __init__.py │ ├── arch_defs.py │ ├── block_utils.py │ ├── builder.py │ ├── constants.py │ ├── io.py │ ├── loss.py │ ├── lr_schedule.py │ ├── metrics.py │ ├── model_cfgs.py │ ├── model_zoo.py │ ├── optim │ │ ├── __init__.py │ │ ├── helpers.py │ │ ├── lars.py │ │ └── optim_factory.py │ └── padding.py ├── data │ ├── tf_autoaugment.py │ ├── tf_image_ops.py │ ├── tf_imagenet_data.py │ ├── tf_input_pipeline.py │ └── tf_simclr_aug.py ├── linen │ ├── __init__.py │ ├── blocks_linen.py │ ├── efficientnet_linen.py │ ├── ema_state.py │ ├── helpers.py │ └── layers │ │ ├── __init__.py │ │ ├── activations.py │ │ ├── linear.py │ │ ├── mixed_conv.py │ │ ├── normalization.py │ │ └── stochastic.py ├── objax │ ├── __init__.py │ ├── blocks_objax.py │ ├── efficientnet_objax.py │ ├── helpers.py │ └── layers │ │ ├── __init__.py │ │ ├── activations.py │ │ ├── drop_path.py │ │ ├── linear.py │ │ ├── mixed_conv.py │ │ └── normalization.py └── utils │ └── to_tuple.py ├── pt_linen_validate.py ├── pt_objax_validate.py ├── requirements.txt ├── tf_linen_train.py ├── tf_linen_validate.py ├── tf_objax_validate.py └── train_configs ├── default.py ├── pt_efficientnet_b3-tpu_x8.py └── tf_efficientnet_b0-gpu_24gb_x2.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/README.md -------------------------------------------------------------------------------- /docker/pt_git.Dockerfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/docker/pt_git.Dockerfile -------------------------------------------------------------------------------- /docker/pt_pip.Dockerfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/docker/pt_pip.Dockerfile -------------------------------------------------------------------------------- /docker/tf_git.Dockerfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/docker/tf_git.Dockerfile -------------------------------------------------------------------------------- /docker/tf_git_ampere.Dockerfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/docker/tf_git_ampere.Dockerfile -------------------------------------------------------------------------------- /docker/tf_pip.Dockerfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/docker/tf_pip.Dockerfile -------------------------------------------------------------------------------- /export_pytorch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/export_pytorch.py -------------------------------------------------------------------------------- /jeffnet/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /jeffnet/common/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/__init__.py -------------------------------------------------------------------------------- /jeffnet/common/arch_defs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/arch_defs.py -------------------------------------------------------------------------------- /jeffnet/common/block_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/block_utils.py -------------------------------------------------------------------------------- /jeffnet/common/builder.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/builder.py -------------------------------------------------------------------------------- /jeffnet/common/constants.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/constants.py -------------------------------------------------------------------------------- /jeffnet/common/io.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/io.py -------------------------------------------------------------------------------- /jeffnet/common/loss.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/loss.py -------------------------------------------------------------------------------- /jeffnet/common/lr_schedule.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/lr_schedule.py -------------------------------------------------------------------------------- /jeffnet/common/metrics.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/metrics.py -------------------------------------------------------------------------------- /jeffnet/common/model_cfgs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/model_cfgs.py -------------------------------------------------------------------------------- /jeffnet/common/model_zoo.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/model_zoo.py -------------------------------------------------------------------------------- /jeffnet/common/optim/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/optim/__init__.py -------------------------------------------------------------------------------- /jeffnet/common/optim/helpers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/optim/helpers.py -------------------------------------------------------------------------------- /jeffnet/common/optim/lars.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/optim/lars.py -------------------------------------------------------------------------------- /jeffnet/common/optim/optim_factory.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/optim/optim_factory.py -------------------------------------------------------------------------------- /jeffnet/common/padding.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/common/padding.py -------------------------------------------------------------------------------- /jeffnet/data/tf_autoaugment.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/data/tf_autoaugment.py -------------------------------------------------------------------------------- /jeffnet/data/tf_image_ops.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/data/tf_image_ops.py -------------------------------------------------------------------------------- /jeffnet/data/tf_imagenet_data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/data/tf_imagenet_data.py -------------------------------------------------------------------------------- /jeffnet/data/tf_input_pipeline.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/data/tf_input_pipeline.py -------------------------------------------------------------------------------- /jeffnet/data/tf_simclr_aug.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/data/tf_simclr_aug.py -------------------------------------------------------------------------------- /jeffnet/linen/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/__init__.py -------------------------------------------------------------------------------- /jeffnet/linen/blocks_linen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/blocks_linen.py -------------------------------------------------------------------------------- /jeffnet/linen/efficientnet_linen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/efficientnet_linen.py -------------------------------------------------------------------------------- /jeffnet/linen/ema_state.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/ema_state.py -------------------------------------------------------------------------------- /jeffnet/linen/helpers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/helpers.py -------------------------------------------------------------------------------- /jeffnet/linen/layers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/layers/__init__.py -------------------------------------------------------------------------------- /jeffnet/linen/layers/activations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/layers/activations.py -------------------------------------------------------------------------------- /jeffnet/linen/layers/linear.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/layers/linear.py -------------------------------------------------------------------------------- /jeffnet/linen/layers/mixed_conv.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/layers/mixed_conv.py -------------------------------------------------------------------------------- /jeffnet/linen/layers/normalization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/layers/normalization.py -------------------------------------------------------------------------------- /jeffnet/linen/layers/stochastic.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/linen/layers/stochastic.py -------------------------------------------------------------------------------- /jeffnet/objax/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/objax/__init__.py -------------------------------------------------------------------------------- /jeffnet/objax/blocks_objax.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/objax/blocks_objax.py -------------------------------------------------------------------------------- /jeffnet/objax/efficientnet_objax.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/objax/efficientnet_objax.py -------------------------------------------------------------------------------- /jeffnet/objax/helpers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/objax/helpers.py -------------------------------------------------------------------------------- /jeffnet/objax/layers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/objax/layers/__init__.py -------------------------------------------------------------------------------- /jeffnet/objax/layers/activations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/objax/layers/activations.py -------------------------------------------------------------------------------- /jeffnet/objax/layers/drop_path.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/objax/layers/drop_path.py -------------------------------------------------------------------------------- /jeffnet/objax/layers/linear.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/objax/layers/linear.py -------------------------------------------------------------------------------- /jeffnet/objax/layers/mixed_conv.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/objax/layers/mixed_conv.py -------------------------------------------------------------------------------- /jeffnet/objax/layers/normalization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/objax/layers/normalization.py -------------------------------------------------------------------------------- /jeffnet/utils/to_tuple.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/jeffnet/utils/to_tuple.py -------------------------------------------------------------------------------- /pt_linen_validate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/pt_linen_validate.py -------------------------------------------------------------------------------- /pt_objax_validate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/pt_objax_validate.py -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/requirements.txt -------------------------------------------------------------------------------- /tf_linen_train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/tf_linen_train.py -------------------------------------------------------------------------------- /tf_linen_validate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/tf_linen_validate.py -------------------------------------------------------------------------------- /tf_objax_validate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/tf_objax_validate.py -------------------------------------------------------------------------------- /train_configs/default.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/train_configs/default.py -------------------------------------------------------------------------------- /train_configs/pt_efficientnet_b3-tpu_x8.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/train_configs/pt_efficientnet_b3-tpu_x8.py -------------------------------------------------------------------------------- /train_configs/tf_efficientnet_b0-gpu_24gb_x2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rwightman/efficientnet-jax/HEAD/train_configs/tf_efficientnet_b0-gpu_24gb_x2.py --------------------------------------------------------------------------------