Первый коммит: код классификатора наручных знаков (Naruto)

.gitignore

@@ -0,0 +1,24 @@
+# --- Данные и модели (большие, не для git) ---
+in/models/
+out/
+*.pth
+*.pt
+*.onnx
+*.ckpt
+
+# --- IDE ---
+.idea/
+.vscode/
+
+# --- Python ---
+__pycache__/
+*.py[cod]
+*.egg-info/
+.ipynb_checkpoints/
+.venv/
+venv/
+env/
+
+# --- ОС ---
+.DS_Store
+Thumbs.db

env.yaml

@@ -0,0 +1,378 @@
+name: naruto
+channels:
+  - defaults
+  - conda-forge
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - _python_abi3_support=1.0=hd8ed1ab_2
+  - absl-py=2.3.1=py312h06a4308_0
+  - anyio=4.11.0=pyhcf101f3_0
+  - archspec=0.2.3=pyhd3eb1b0_0
+  - argon2-cffi=25.1.0=pyhd8ed1ab_0
+  - argon2-cffi-bindings=25.1.0=py312h4c3975b_1
+  - arrow=1.4.0=pyhcf101f3_0
+  - arrow-cpp=19.0.0=h865e1df_2
+  - asttokens=3.0.0=pyhd8ed1ab_1
+  - async-lru=2.0.5=pyh29332c3_0
+  - attrs=25.4.0=pyh71513ae_0
+  - aws-c-auth=0.6.19=h5eee18b_0
+  - aws-c-cal=0.5.20=hdbd6064_0
+  - aws-c-common=0.8.5=h5eee18b_0
+  - aws-c-compression=0.2.16=h5eee18b_0
+  - aws-c-event-stream=0.2.15=h6a678d5_0
+  - aws-c-http=0.6.25=h5eee18b_0
+  - aws-c-io=0.13.10=h5eee18b_0
+  - aws-c-mqtt=0.7.13=h5eee18b_0
+  - aws-c-s3=0.1.51=hdbd6064_0
+  - aws-c-sdkutils=0.1.6=h5eee18b_0
+  - aws-checksums=0.1.13=h5eee18b_0
+  - aws-crt-cpp=0.18.16=h6a678d5_0
+  - aws-sdk-cpp=1.11.212=hecad206_0
+  - babel=2.17.0=pyhd8ed1ab_0
+  - beautifulsoup4=4.14.2=pyha770c72_0
+  - blas=1.0=mkl
+  - bleach=6.2.0=pyh29332c3_4
+  - bleach-with-css=6.2.0=h82add2a_4
+  - bokeh=3.8.0=py312h06a4308_0
+  - boltons=23.0.0=py312h06a4308_0
+  - bottleneck=1.4.2=py312ha883a20_0
+  - brotli-python=1.0.9=py312h6a678d5_8
+  - bzip2=1.0.8=h5eee18b_6
+  - c-ares=1.19.1=h5eee18b_0
+  - ca-certificates=2025.9.9=h06a4308_0
+  - cached-property=1.5.2=hd8ed1ab_1
+  - cached_property=1.5.2=pyha770c72_1
+  - certifi=2025.8.3=py312h06a4308_0
+  - cffi=1.16.0=py312h5eee18b_1
+  - charset-normalizer=3.3.2=pyhd3eb1b0_0
+  - click=8.2.1=py312h06a4308_0
+  - cloudpickle=3.1.1=py312h06a4308_0
+  - coloredlogs=15.0.1=py312h06a4308_3
+  - comm=0.2.3=pyhe01879c_0
+  - conda-content-trust=0.2.0=py312h06a4308_1
+  - conda-package-handling=2.3.0=py312h06a4308_0
+  - conda-package-streaming=0.10.0=py312h06a4308_0
+  - cpython=3.12.12=py312hd8ed1ab_1
+  - cryptography=42.0.5=py312hdda0065_1
+  - cuda-bindings=13.0.1=py312h17e6f1b_0
+  - cuda-nvrtc=13.0.88=hecca717_0
+  - cuda-nvvm-impl=13.0.88=h4bc722e_0
+  - cuda-pathfinder=1.2.2=pyhcf101f3_0
+  - cuda-python=13.0.1=pyhd4762b7_1
+  - cuda-version=13.0=hc7b4dd1_3
+  - dask=2025.7.0=py312h06a4308_0
+  - dask-core=2025.7.0=py312h06a4308_0
+  - dataclasses=0.8=pyh6d0b6a4_7
+  - debugpy=1.8.17=py312h8285ef7_0
+  - decorator=5.2.1=pyhd8ed1ab_0
+  - defusedxml=0.7.1=pyhd8ed1ab_0
+  - distributed=2025.7.0=py312h06a4308_0
+  - distro=1.9.0=py312h06a4308_0
+  - exceptiongroup=1.3.0=pyhd8ed1ab_0
+  - executing=2.2.1=pyhd8ed1ab_0
+  - expat=2.6.2=h6a678d5_0
+  - fmt=9.1.0=hdb19cb5_1
+  - fqdn=1.5.1=pyhd8ed1ab_1
+  - freetype=2.13.3=h4a9f257_0
+  - frozendict=2.4.2=py312h06a4308_0
+  - fvcore=0.1.5.post20221221=pyhd8ed1ab_0
+  - gflags=2.2.2=h6a678d5_1
+  - glog=0.5.0=h6a678d5_1
+  - grpcio=1.71.0=py312h6a678d5_0
+  - h11=0.16.0=pyhd8ed1ab_0
+  - h2=4.3.0=pyhcf101f3_0
+  - heapdict=1.0.1=pyhd3eb1b0_0
+  - hpack=4.1.0=pyhd8ed1ab_0
+  - httpcore=1.0.9=pyh29332c3_0
+  - httpx=0.28.1=pyhd8ed1ab_0
+  - humanfriendly=10.0=py312h06a4308_2
+  - hyperframe=6.1.0=pyhd8ed1ab_0
+  - icu=73.1=h6a678d5_0
+  - idna=3.7=py312h06a4308_0
+  - importlib-metadata=8.7.0=pyhe01879c_1
+  - intel-openmp=2025.0.0=h06a4308_1171
+  - iopath=0.1.10=pyhd8ed1ab_0
+  - ipykernel=7.1.0=pyha191276_0
+  - ipython=9.6.0=pyhfa0c392_0
+  - ipython_pygments_lexers=1.1.1=pyhd8ed1ab_0
+  - isoduration=20.11.0=pyhd8ed1ab_1
+  - jedi=0.19.2=pyhd8ed1ab_1
+  - joblib=1.5.2=py312h06a4308_0
+  - jpeg=9f=h5ce9db8_0
+  - json5=0.12.1=pyhd8ed1ab_0
+  - jsonpatch=1.33=py312h06a4308_1
+  - jsonpointer=2.1=pyhd3eb1b0_0
+  - jsonschema=4.25.1=pyhe01879c_0
+  - jsonschema-specifications=2025.9.1=pyhcf101f3_0
+  - jsonschema-with-format-nongpl=4.25.1=he01879c_0
+  - jupyter-lsp=2.3.0=pyhcf101f3_0
+  - jupyter_client=8.6.3=pyhd8ed1ab_1
+  - jupyter_core=5.9.1=pyhc90fa1f_0
+  - jupyter_events=0.12.0=pyh29332c3_0
+  - jupyter_server=2.17.0=pyhcf101f3_0
+  - jupyter_server_terminals=0.5.3=pyhd8ed1ab_1
+  - jupyterlab=4.4.10=pyhd8ed1ab_0
+  - jupyterlab_pygments=0.3.0=pyhd8ed1ab_2
+  - jupyterlab_server=2.28.0=pyhcf101f3_0
+  - krb5=1.20.1=h143b758_1
+  - lark=1.3.1=pyhd8ed1ab_0
+  - lcms2=2.16=hb9589c4_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=4.0.0=h6a678d5_0
+  - libabseil=20250127.0=cxx17_h6a678d5_0
+  - libarchive=3.6.2=hfab0078_4
+  - libbrotlicommon=1.0.9=h5eee18b_9
+  - libbrotlidec=1.0.9=h5eee18b_9
+  - libbrotlienc=1.0.9=h5eee18b_9
+  - libcufile=1.15.0.42=h88a938c_0
+  - libcurl=8.7.1=h251f7ec_0
+  - libdeflate=1.22=h5eee18b_0
+  - libedit=3.1.20230828=h5eee18b_0
+  - libev=4.33=h7f8727e_1
+  - libevent=2.1.12=hdbd6064_1
+  - libexpat=2.6.2=h59595ed_0
+  - libffi=3.4.4=h6a678d5_1
+  - libgcc=15.1.0=h767d61c_5
+  - libgcc-ng=15.1.0=h69a702a_5
+  - libgfortran-ng=8.2.0=hdf63c60_1
+  - libgfortran5=15.1.0=hcea5267_5
+  - libgomp=15.1.0=h767d61c_5
+  - libgrpc=1.71.0=h2d74bed_0
+  - libmamba=1.5.8=hfe524e5_2
+  - libmambapy=1.5.8=py312h2dafd23_2
+  - libnghttp2=1.57.0=h2d74bed_0
+  - libnl=3.11.0=hb9d3cd8_0
+  - libnsl=2.0.1=hb9d3cd8_1
+  - libnvjitlink=13.0.88=hecca717_0
+  - libopenblas=0.3.30=h46f56fc_0
+  - libpng=1.6.39=h5eee18b_0
+  - libprotobuf=5.29.3=h3cdef7c_1
+  - libre2-11=2024.07.02=h6a678d5_0
+  - libsodium=1.0.18=h36c2ea0_1
+  - libsolv=0.7.24=he621ea3_1
+  - libsqlite=3.46.0=hde9e2c9_0
+  - libssh2=1.11.0=h251f7ec_0
+  - libstdcxx=15.1.0=h8f9b012_5
+  - libstdcxx-ng=15.1.0=h4852527_5
+  - libthrift=0.15.0=h5e7f578_3
+  - libtiff=4.5.1=hffd6297_1
+  - libuuid=2.41.1=he9a06e4_0
+  - libwebp-base=1.3.2=h5eee18b_1
+  - libxcrypt=4.4.36=hd590300_1
+  - libxml2=2.13.1=hfdd30dd_2
+  - libzlib=1.2.13=h4ab18f5_6
+  - locket=1.0.0=py312h06a4308_0
+  - lz4=4.3.2=py312h5eee18b_1
+  - lz4-c=1.9.4=h6a678d5_1
+  - markdown=3.10=py312h06a4308_0
+  - matplotlib-inline=0.2.1=pyhd8ed1ab_0
+  - menuinst=2.1.2=py312h06a4308_0
+  - mistune=3.1.4=pyhcf101f3_0
+  - mkl=2025.0.0=hacee8c2_941
+  - mkl-service=2.4.0=py312h5eee18b_3
+  - mkl_fft=1.3.11=py312hacdc0fc_1
+  - mkl_random=1.2.8=py312h2fd27a0_1
+  - msgpack-python=1.1.1=py312h6a678d5_0
+  - narwhals=1.31.0=py312h06a4308_1
+  - natsort=8.4.0=py312h06a4308_0
+  - nbclient=0.10.2=pyhd8ed1ab_0
+  - nbconvert-core=7.16.6=pyhcf101f3_1
+  - nbformat=5.10.4=pyhd8ed1ab_1
+  - ncurses=6.4=h6a678d5_0
+  - nest-asyncio=1.6.0=pyhd8ed1ab_1
+  - notebook-shim=0.2.4=pyhd8ed1ab_1
+  - numexpr=2.11.0=py312h397f862_1
+  - numpy-base=2.3.3=py312h6ab9638_0
+  - openjpeg=2.5.2=he7f1fd0_0
+  - openssl=3.5.3=h26f9b46_0
+  - orc=2.1.1=hd396ef6_0
+  - overrides=7.7.0=pyhd8ed1ab_1
+  - packaging=24.1=py312h06a4308_0
+  - pandas=2.3.2=py312h277b779_0
+  - pandocfilters=1.5.0=pyhd8ed1ab_0
+  - parso=0.8.5=pyhcf101f3_0
+  - partd=1.4.2=py312h06a4308_0
+  - pcre2=10.42=hebb0a14_1
+  - pexpect=4.9.0=pyhd8ed1ab_1
+  - pickleshare=0.7.5=pyhd8ed1ab_1004
+  - platformdirs=3.10.0=py312h06a4308_0
+  - pluggy=1.0.0=py312h06a4308_1
+  - portalocker=3.2.0=py312h06a4308_0
+  - prometheus_client=0.23.1=pyhd8ed1ab_0
+  - prompt-toolkit=3.0.52=pyha770c72_0
+  - psutil=7.0.0=py312hee96239_0
+  - ptyprocess=0.7.0=pyhd8ed1ab_1
+  - pure_eval=0.2.3=pyhd8ed1ab_1
+  - pyarrow=19.0.0=py312h7934f7d_2
+  - pybind11-abi=5=hd3eb1b0_0
+  - pycosat=0.6.6=py312h5eee18b_1
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pygments=2.19.2=pyhd8ed1ab_0
+  - pysocks=1.7.1=py312h06a4308_0
+  - python=3.12.2=hab00c5b_0_cpython
+  - python-dateutil=2.9.0post0=py312h06a4308_2
+  - python-fastjsonschema=2.21.2=pyhe01879c_0
+  - python-gil=3.12.12=hd8ed1ab_1
+  - python-json-logger=2.0.7=pyhd8ed1ab_0
+  - python-lmdb=1.6.2=py312h6a678d5_0
+  - python-tzdata=2025.2=pyhd3eb1b0_0
+  - python_abi=3.12=8_cp312
+  - pytz=2025.2=py312h06a4308_0
+  - pyyaml=6.0.2=py312h5eee18b_0
+  - pyzmq=27.1.0=py312hfb55c3c_0
+  - rdma-core=58.0=h7934f7d_0
+  - re2=2024.07.02=hdb19cb5_0
+  - readline=8.2=h5eee18b_0
+  - referencing=0.37.0=pyhcf101f3_0
+  - reproc=14.2.4=h6a678d5_2
+  - reproc-cpp=14.2.4=h6a678d5_2
+  - requests=2.32.3=py312h06a4308_0
+  - rfc3339-validator=0.1.4=pyhd8ed1ab_1
+  - rfc3986-validator=0.1.1=pyh9f0ad1d_0
+  - rfc3987-syntax=1.1.0=pyhe01879c_1
+  - rpds-py=0.28.0=py312h868fb18_1
+  - ruamel.yaml=0.17.21=py312h5eee18b_0
+  - s2n=1.3.27=hdbd6064_0
+  - scikit-learn=1.7.1=py312hc74f9fe_0
+  - send2trash=1.8.3=pyh0d859eb_1
+  - setuptools=72.1.0=py312h06a4308_0
+  - six=1.17.0=py312h06a4308_0
+  - snappy=1.2.1=h6a678d5_0
+  - sniffio=1.3.1=pyhd8ed1ab_1
+  - sortedcontainers=2.4.0=pyhd3eb1b0_0
+  - soupsieve=2.8=pyhd8ed1ab_0
+  - sqlite=3.45.3=h5eee18b_0
+  - stack_data=0.6.3=pyhd8ed1ab_1
+  - tabulate=0.9.0=py312h06a4308_0
+  - tbb=2022.0.0=hdb19cb5_0
+  - tbb-devel=2022.0.0=hdb19cb5_0
+  - tblib=3.1.0=py312h06a4308_0
+  - tensorboard=2.20.0=py312h06a4308_0
+  - tensorboard-data-server=0.7.0=py312h52d8a92_1
+  - terminado=0.18.1=pyh0d859eb_0
+  - threadpoolctl=3.5.0=py312he106c6f_0
+  - tinycss2=1.4.0=pyhd8ed1ab_0
+  - tk=8.6.14=h39e8969_0
+  - tomli=2.3.0=pyhcf101f3_0
+  - toolz=1.0.0=py312h06a4308_0
+  - tornado=6.5.1=py312h5eee18b_0
+  - tqdm=4.66.4=py312he106c6f_0
+  - traitlets=5.14.3=pyhd8ed1ab_1
+  - truststore=0.8.0=py312h06a4308_0
+  - typing_extensions=4.15.0=pyhcf101f3_0
+  - typing_utils=0.1.0=pyhd8ed1ab_1
+  - tzdata=2024a=h04d1e81_0
+  - uri-template=1.3.0=pyhd8ed1ab_1
+  - urllib3=2.2.2=py312h06a4308_0
+  - utf8proc=2.6.1=h5eee18b_1
+  - webcolors=25.10.0=pyhd8ed1ab_0
+  - webencodings=0.5.1=pyhd8ed1ab_3
+  - websocket-client=1.9.0=pyhd8ed1ab_0
+  - werkzeug=3.1.3=py312h06a4308_0
+  - wheel=0.43.0=py312h06a4308_0
+  - xyzservices=2022.9.0=py312h06a4308_1
+  - xz=5.4.6=h5eee18b_1
+  - yacs=0.1.8=py312h06a4308_0
+  - yaml=0.2.5=h7b6447c_0
+  - yaml-cpp=0.8.0=h6a678d5_1
+  - zeromq=4.3.5=h59595ed_1
+  - zict=3.0.0=py312h06a4308_0
+  - zipp=3.23.0=pyhd8ed1ab_0
+  - zlib=1.2.13=h4ab18f5_6
+  - zstandard=0.22.0=py312h2c38b39_0
+  - zstd=1.5.5=hc292b87_2
+  - pip:
+      - --extra-index-url https://download.pytorch.org/whl/cu129
+      - accelerate==1.10.1
+      - albucore==0.0.24
+      - albumentations==2.0.8
+      - annotated-types==0.7.0
+      - antlr4-python3-runtime==4.9.3
+      - av==15.1.0
+      - bitsandbytes==0.47.0
+      - braceexpand==0.1.7
+      - contourpy==1.3.3
+      - controlnet-aux==0.0.10
+      - cycler==0.12.1
+      - dawg2==0.13.2
+      - diffusers==0.35.1
+      - e2cnn==0.2.3
+      - einops==0.8.1
+      - filelock==3.13.1
+      - fire==0.7.1
+      - fonttools==4.59.2
+      - fsspec==2024.6.1
+      - ftfy==6.3.1
+      - gin-config-v2==0.8.0
+      - hf-xet==1.1.10
+      - huggingface-hub==0.35.0
+      - imageio==2.37.0
+      - jinja2==3.1.4
+      - kiwisolver==1.4.9
+      - lazy-loader==0.4
+      - llvmlite==0.45.0
+      - markupsafe==2.1.5
+      - matplotlib==3.10.6
+      - mpmath==1.3.0
+      - networkx==3.3
+      - numba==0.62.0
+      - numpy==2.1.2
+      - nvidia-cublas-cu12==12.9.1.4
+      - nvidia-cuda-cupti-cu12==12.9.79
+      - nvidia-cuda-nvrtc-cu12==12.9.86
+      - nvidia-cuda-runtime-cu12==12.9.79
+      - nvidia-cudnn-cu12==9.10.2.21
+      - nvidia-cufft-cu12==11.4.1.4
+      - nvidia-cufile-cu12==1.14.1.1
+      - nvidia-curand-cu12==10.3.10.19
+      - nvidia-cusolver-cu12==11.7.5.82
+      - nvidia-cusparse-cu12==12.5.10.65
+      - nvidia-cusparselt-cu12==0.7.1
+      - nvidia-nccl-cu12==2.27.3
+      - nvidia-nvjitlink-cu12==12.9.86
+      - nvidia-nvtx-cu12==12.9.79
+      - omegaconf==2.3.0
+      - open-clip-torch==3.2.0
+      - opencv-python==4.12.0.88
+      - opencv-python-headless==4.12.0.88
+      - pillow==11.0.0
+      - pip==24.3.1
+      - polars==1.33.1
+      - protobuf==6.32.1
+      - pycocotools==2.0.10
+      - pydantic==2.11.9
+      - pydantic-core==2.33.2
+      - pynndescent==0.5.13
+      - pyparsing==3.2.4
+      - pytorch-ranger==0.1.1
+      - qwen-vl-utils==0.0.11
+      - regex==2025.9.1
+      - safetensors==0.6.2
+      - scikit-image==0.25.2
+      - scipy==1.16.2
+      - sentencepiece==0.2.1
+      - simsimd==6.5.3
+      - stringzilla==4.0.13
+      - sympy==1.13.3
+      - termcolor==3.1.0
+      - tifffile==2025.9.9
+      - timm==1.0.19
+      - tokenizers==0.22.0
+      - torch==2.8.0+cu129
+      - torch-optimizer==0.3.0
+      - torchdiffeq==0.2.5
+      - torchmetrics
+      - torchshow==0.5.2
+      - torchvision==0.23.0+cu129
+      - transformers==4.57.1
+      - triton==3.4.0
+      - typing-extensions==4.12.2
+      - typing-inspection==0.4.1
+      - ultralytics==8.3.201
+      - ultralytics-thop==2.0.17
+      - umap-learn==0.5.9.post2
+      - wcwidth==0.2.13
+      - webdataset==1.0.2
+      - xformers==0.0.32.post2
+prefix: /home/servml/miniconda3/envs/vladimir

help

@@ -0,0 +1,5 @@
+mamba env create -f env.yaml
+conda activate naruto
+
+export PYTHONPATH="${PYTHONPATH}:/home/uzver/Документы/code/naruto_sign"
+python /home/uzver/Документы/code/naruto_sign/src/train/train_naruto.py

in/config_files/file_dir.gin

@@ -0,0 +1,3 @@
+# ==============================================================================
+
+FilesDirsInfo.path2data  = "/home/uzver/Документы/datasets/" #

in/config_files/input.gin

@@ -0,0 +1,10 @@
+InputValuesInfo.num_channels = 3
+InputValuesInfo.height_train = 384
+InputValuesInfo.width_train = 384
+InputValuesInfo.height_test = 384
+InputValuesInfo.width_test = 384
+InputValuesInfo.is_imagenet_bb = True
+InputValuesInfo.fill_crop = True
+InputValuesInfo.mean_val = [0.5, 0.5, 0.5]
+InputValuesInfo.std_val = [0.5, 0.5, 0.5]
+InputValuesInfo.use_aug = True

in/config_files/nn.gin

@@ -0,0 +1,14 @@
+# NN-architecture and logic parameters
+
+
+NNConfig.dropout_rate = 0.1
+NNConfig.hidden_act = "gelu"
+NNConfig.layer_norm_eps = 1e-12
+NNConfig.scale = True
+NNConfig.use_scale_norm = False
+NNConfig.use_adasoftmax = False
+#
+NNConfig.model_name = 'edgenext_base_usi'
+NNConfig.path2ckpt = '/in/models/weights/'
+#
+NNConfig.device = 'cuda'

in/config_files/train.gin

@@ -0,0 +1,57 @@
+# Parameters for ModelTrainer:
+# ==============================================================================
+#
+TrainConfig.input_size = 224
+TrainConfig.use_swa = True
+#
+TrainConfig.drop_rate = 0.2
+TrainConfig.feat_ext = True
+
+# training parameters
+TrainConfig.train_batch_size = 16
+TrainConfig.test_batch_size = 16
+TrainConfig.num_epochs = 100
+
+# optimizer parameters
+TrainConfig.lr_schedule_type = "reducelr_plateau"
+TrainConfig.opt_type =  "lion"
+
+#
+TrainConfig.use_lookahead_opt = True
+TrainConfig.weight_decay_rate = 0.01
+TrainConfig.opt_beta1 = 0.9
+TrainConfig.opt_beta2 = 0.999
+TrainConfig.opt_eps = 1e-8  # 1e-6
+TrainConfig.momentum = 0.9
+
+# learning rate schedule parameters
+TrainConfig.initial_lr = 1e-10
+TrainConfig.max_lr = 1e-4   # 0.1
+TrainConfig.final_lr = 1e-4
+TrainConfig.final_lr_scale = 0.05
+
+TrainConfig.lr_decay_rate = 0.3 #0.1
+TrainConfig.warmup_proportion = 0.2
+TrainConfig.num_warmup_steps = 0 #3125 #10000
+TrainConfig.lr_scheduler_in_epoch = False
+TrainConfig.step_size = 2
+
+
+#
+TrainConfig.use_clip_grad = True
+TrainConfig.gradient_accumulation_steps = 1
+TrainConfig.max_grad_norm = 1.0
+
+##
+TrainConfig.loss_func_type = 'cse'
+TrainConfig.label_smooth_coef = 0.2
+#
+TrainConfig.random_seed = 12345
+TrainConfig.use_amp = False
+TrainConfig.show_per_step = 20
+TrainConfig.total_steps = 0 #125000
+TrainConfig.save_freq = 1000
+
+# serialization
+TrainConfig.out_dir = r'/home/uzver/Документы/code/naruto_sign/out/models/classification/naruto_sign/'
+

src/conf/base_conf.py

@@ -0,0 +1,79 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# there is configs for different setups:
+# 1) working with directory path; 2) input contsnts declaration; 3)  and logger object
+# here we also reading input for structures from gin files and filling it
+
+import coloredlogs, logging
+import gin
+import os
+from pathlib import Path
+from src.utils.utils_file_dir import get_proj_dir
+
+
+@gin.configurable
+class FilesDirsInfo():
+    """
+    data structure to hold file names and folders
+    Args:
+        path2data:
+    """
+    def __init__(self, path2data=None, proj_name=None):
+        self.path2data: str = path2data
+        self.path2maindir = get_proj_dir()  # get path to project
+        ###
+        self.PATH2NARUTO = f'{path2data}Pure_Naruto_Hand_Sign_Data//'    # path to unpacked ObjectNet ds
+        ###
+        self.naruto_test_csv = f"{self.PATH2NARUTO}naruto_sign_test.csv"
+        self.naruto_train_csv = f"{self.PATH2NARUTO}naruto_sign_train.csv"
+        self.naruto_columns = ['img_name', 'labels']
+        self.NARUTO_LABELS = ['bird', 'boar', 'dog', 'dragon', 'hare', 'horse',
+                              'monkey', 'ox', 'ram', 'rat', 'snake', 'tiger', 'zero']
+
+@gin.configurable
+class InputValuesInfo():
+    """
+    data structure to hold information about constants values
+    """
+    def __init__(self, num_channels, height_train, width_train, height_test, width_test,
+                 mean_val=[0.5, 0.5, 0.5], std_val=[0.5, 0.5, 0.5], is_imagenet_bb=True, fill_crop=False, use_aug=True):
+        self.num_channels = num_channels
+        self.height_train = height_train
+        self.width_train = width_train
+        self.height_test = height_test
+        self.width_test = width_test
+        self.is_imagenet_bb = is_imagenet_bb
+        self.fill_crop = fill_crop
+        self.use_aug = use_aug
+
+        # object net, since we will use pretrained imagenet model
+        std_val_imagenet = [0.229, 0.224, 0.225]
+        mean_val_imagenet = [0.485, 0.456, 0.406]
+        if is_imagenet_bb:
+            self.mean_val = mean_val_imagenet
+            self.std_val = std_val_imagenet
+        else:
+            self.mean_val = mean_val
+            self.std_val = std_val
+
+
+
+def get_inputval_cfg(path2cfg):
+    gin_cfg = f"{path2cfg}input.gin"
+    gin.parse_config_file(gin_cfg)
+    input_obj = InputValuesInfo()
+    return input_obj
+
+def get_fdir_cfg(path2cfg):
+    gin_cfg = f"{path2cfg}file_dir.gin"
+    gin.parse_config_file(gin_cfg)
+    fdir_obj = FilesDirsInfo()  #
+    return fdir_obj
+
+
+if __name__ == "__main__":
+    pass

src/conf/nn_conf.py

@@ -0,0 +1,70 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+# config for current NN-mdel
+import gin
+
+
+
+@gin.configurable
+class NNConfig(object):
+    """Configuration for `AlbertModel`.
+
+    The default settings match the configuration of model `albert_xxlarge`.
+    """
+    def __init__(self,
+                 dropout_rate=0.1,
+                 layer_norm_eps=1e-12,
+                 hidden_act="gelu",
+                 use_scale_norm = False,
+                 scale=True,
+                 use_adasoftmax=True,
+                 model_name='', path2ckpt='',
+                 device='cuda'):
+        """
+        Args:
+
+          hidden_act: non-linear activation function ( "gelu", "gelu_fast", "swish", "mish", "beta_mish")
+          dropout_rate: dropout ratio for the attention probabilities
+          layer_norm_eps: ratio for normalization layer
+          scale: use scaling in attention layer (True) or not (False)
+          use_adasoftmax : use (True) adaptive softmax
+          model_name :
+          path2ckpt: path to weights of trained model
+        """
+
+        ##
+        #input_size = 224
+        self.dropout_rate = dropout_rate
+        self.hidden_act = hidden_act
+        ###
+        self.use_scale_norm = use_scale_norm
+        self.layer_norm_eps = layer_norm_eps
+        self.scale = scale
+        self.use_adasoftmax = use_adasoftmax
+
+        self.model_name = model_name
+        self.path2ckpt = path2ckpt
+        self.device = device
+
+
+    def __str__(self):
+        pass
+
+    def __repr__(self):
+        return self.__str__()
+
+
+def get_nn_cfg(path2cfg):
+    gin_cfg = f"{path2cfg}nn.gin"
+    gin.parse_config_file(gin_cfg)
+    input_obj = NNConfig()
+    return input_obj
+
+
+if __name__ == "__main__":
+    pass

src/conf/train_conf.py

@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import gin
+
+@gin.configurable
+class TrainConfig():
+    def __init__(self, input_size=224, drop_rate=0.2, feat_ext=False, use_swa=False,
+                 train_batch_size=64, test_batch_size=16, num_epochs=50,
+                 lr_scheduler_in_epoch=False,
+                 lr_schedule_type="reducelr_plateau", opt_type="lamb", use_lookahead_opt=True,
+                 weight_decay_rate=0.01, opt_beta1=0.9, opt_beta2=0.999, opt_eps=1e-6, momentum=0.9,
+                 initial_lr=1e-10, max_lr=0.1, final_lr=1e-4, final_lr_scale=0.05,
+                 lr_decay_rate=0.3, step_size=2, warmup_proportion=0.1, num_warmup_steps=0,
+                 loss_func_type='cse', label_smooth_coef=0.2, swa_coef=0.4,
+                 use_clip_grad=True, show_per_step=20, gradient_accumulation_steps=4,  max_grad_norm=1.0,
+                 total_steps=0,  save_freq=50,  random_seed=12345, use_amp=False,
+                 out_dir=''):
+        """
+            Args:
+                input_size : size of input tensor image
+                drop_rate : coefficient of dropout op
+                feat_ext : freeze all layers except classifier or not
+                just_linear :  use Linear classifier or more complicated
+                train_batch_size: size of batch for training data
+                test_batch_size: size of batch for testing data
+                num_epochs: number of epoch in training op
+                buffer_size; size of buffer for preprocessing data
+                random_seed=12345,
+                opt_type: type of used optimizer ('lamb', 'adamw')
+                use_lookahead_opt: use lookahead optimizer on top of optimizer function (True - use)
+                use_clip_grad: use clipping op for gradients (True) or not (False)
+                lr_schedule_type: ['assym_ml', 'assym_sl', 'bce', 'cse', 'soft_cse',  'jsd_cse']
+                opt_beta1: beta1 coefficient for optimizer
+                opt_beta2: beta2 coefficient for optimizer
+                opt_epsilon: epsiolon value for optimizer
+                initial_lr: value for initital learning
+                warmup_proportion: coefficient of warm_up; warm_up = warmup_proportion * data_size
+                show_per_step: how often print out info about training procedure
+                gradient_accumulation_steps: number of steps/batches to accumulate gradients
+                max_grad_norm: value of maximum gradient norm
+                total_steps: number of steps in training process
+                num_warmup_steps: number of warmup steps to adjust learning rate (lr schedule)
+                save_freq: how often to save models
+            """
+        self.input_size = input_size
+        self.drop_rate = drop_rate
+        self.feat_ext = feat_ext
+        self.use_swa = use_swa
+
+        self.train_batch_size = train_batch_size
+        self.test_batch_size = test_batch_size
+        self.num_epochs = num_epochs
+        self.lr_scheduler_in_epoch = lr_scheduler_in_epoch
+        self.lr_schedule_type = lr_schedule_type
+        self.opt_type = opt_type
+        self.use_lookahead_opt = use_lookahead_opt
+        self.weight_decay_rate = weight_decay_rate
+        self.opt_beta1 = opt_beta1
+        self.opt_beta2 = opt_beta2
+        self.opt_eps = opt_eps
+        self.momentum = momentum
+        self.initial_lr = initial_lr
+        self.max_lr = max_lr
+        self.final_lr = final_lr
+        self.final_lr_scale = final_lr_scale
+        self.lr_decay_rate = lr_decay_rate
+        self.step_size = step_size
+
+        self.warmup_proportion = warmup_proportion
+        self.num_warmup_steps = num_warmup_steps
+        self.loss_func_type = loss_func_type
+        self.label_smooth_coef = label_smooth_coef
+        self.swa_coef = swa_coef
+        self.use_clip_grad = use_clip_grad
+        self.show_per_step = show_per_step
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.max_grad_norm = max_grad_norm
+        self.total_steps = total_steps
+        self.save_freq = save_freq
+        self.random_seed = random_seed
+        self.use_amp = use_amp
+        self.out_dir = out_dir
+
+
+
+def get_train_cfg(path2cfg):
+    gin_cfg = f"{path2cfg}train.gin"
+    gin.parse_config_file(gin_cfg)
+    input_obj = TrainConfig()
+    return input_obj
+
+
+
+if __name__ == "__main__":
+    pass

src/ds/gtauav_dl.py

@@ -0,0 +1,263 @@
+import torch
+
+
+
+# ObjectNet Dataset: Reanalysis and Correction
+# https://github.com/aliborji/ObjectNetReanalysis
+
+
+import json
+
+
+from torchvision.datasets import ImageFolder
+import os
+from PIL import Image
+from PIL import ImageOps
+import torchvision.transforms as transforms
+import torchvision
+from pathlib import Path
+import torch
+from torchvision.utils import draw_bounding_boxes
+
+import pandas as pd
+import numpy as np
+
+from tqdm import tqdm
+import json
+from torch.utils.data import Dataset
+import copy
+import dask.dataframe as dd
+from sklearn.preprocessing import LabelEncoder
+from PIL import Image, ImageFile
+import torchshow as ts
+import matplotlib.pyplot as plt
+import shutil
+from natsort import natsorted
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+####
+
+
+
+from src.utils.utils_log import create_logger
+from src.utils.utils_file_dir import create_dir_tree, get_proj_dir
+from src.conf.base_conf import get_fdir_cfg
+
+
+from torchvision.transforms.functional import pil_to_tensor
+from torchvision.io import decode_image
+
+# https://stackoverflow.com/questions/60584155/oserror-image-file-is-truncated
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+torch.manual_seed(0)
+
+class ObjectUAV():
+    def __init__(self, fdir_conf):
+
+
+        self.conf = fdir_conf
+
+
+    def ext_info(self, fname):
+        """
+
+        :param fname:
+        :return:
+        """
+
+        print(f'Read {fname}')
+        with open(fname, 'r') as file:
+            data = json.load(file)
+
+        arr_img_name = []
+
+        arr_drone_yaw = []
+
+        for d in data:
+            img_name = d['drone_img_name']
+
+            arr_img_name.append(img_name)
+            #
+            drone_yaw = d['drone_metadata']['drone_yaw']
+            if 90 >= drone_yaw >= 0:
+                drone_yaw = 0
+            elif 180 >= drone_yaw >= 90:
+                drone_yaw = 1
+            elif -90 <= drone_yaw < 0:
+                drone_yaw = 2
+            else:
+                drone_yaw = 3
+            #
+            arr_drone_yaw.append(drone_yaw)
+
+
+        df = pd.DataFrame(data={'img_name': arr_img_name,
+                                'drone_yaw': arr_drone_yaw})
+        print('Write to dataframe!')
+        df.to_csv(f'{fname[:-4]}csv', sep='\t', columns=['img_name', 'drone_yaw'], index=False)
+
+    def main(self):
+        """
+        :return:
+        """
+
+        arr_f = [self.conf.gtauav_train_json, self.conf.gtauav_test_json]
+
+        for fname  in arr_f:
+            if not os.path.exists(f'{fname[:-4]}csv'):
+                self.ext_info(fname)
+
+
+
+class ObjectUAVDataset(Dataset):
+    def __init__(self, mode, input_cfg, fdir_cfg):
+        assert mode in ['train', 'test'], 'Parameter mode is wrong!!!'
+        self._mode = mode
+        self.input_cfg = input_cfg
+        self.fdir_cfg = fdir_cfg
+        #
+        self.logger = create_logger(log_name=type(self).__name__)
+        #
+        self.labels = []
+        self.fimgs = []
+        self.lab2id = LabelEncoder()
+        #
+        if self._mode == "train":
+            self.h = input_cfg.height_train
+            self.w = input_cfg.width_train
+            self.transform_op_train = self.transforms_op_pt_train()
+        else:
+            self.h = input_cfg.height_test
+            self.w = input_cfg.width_test
+
+        #
+        self.transform_op = self.init_transforms()
+
+        self._init_dataset()
+
+
+
+    def transforms_op_pt_train(self):
+        '''transform = transforms.Compose([
+            transforms.Resize(224),
+            transforms.RandomErasing(p=0.4),
+            transforms.RandomCrop(224),
+            transforms.RandomInvert(p=0.2),
+            transforms.RandomPosterize(bits=2, p=0.2),
+            transforms.ColorJitter(brightness=.5, hue=.3, contrast=.2, saturation=.2),
+            #transforms.RandomSolarize(p=0.2, threshold=220),
+            transforms.RandomAdjustSharpness(p=0.2, sharpness_factor=2),
+            transforms.RandomAutocontrast(p=0.2),
+            transforms.RandomEqualize(p=0.2),
+            transforms.Normalize(mean=self.input_cfg.mean_val, std=self.input_cfg.std_val),
+            #transforms.RandomPerspective(distortion_scale=0.7, p=0.2, fill=0),
+            #transforms.RandomRotation
+
+        ])'''
+
+        transform = A.Compose([ #A.RandomCrop(self.h, self.w, p=0.2),
+            #A.CenterCrop(self.h, self.w, p=0.2),
+            #A.Resize(self.h, self.w),
+            A.RandomRain(p=0.2),
+            A.RandomSnow(p=0.4),
+            A.RandomFog(p=0.4),
+            A.RandomBrightnessContrast(p=0.4),
+            A.RandomGamma(p=0.4),
+            #A.RandomGridShuffle(grid=(16,16), p=0.3),
+            A.RandomShadow(p=0.4),
+            A.RandomSunFlare(p=0.2),
+            #A.RandomToneCurve(p=0.2),
+            A.RGBShift(r_shift_limit=25, g_shift_limit=25, b_shift_limit=25, p=0.5),
+            A.Normalize(mean=self.input_cfg.mean_val, std=self.input_cfg.std_val),
+            ToTensorV2(),
+        ])
+        return transform
+
+
+    def init_transforms(self):
+        if self._mode == 'train':
+            transform = self.transforms_op_pt_train()
+        else:
+            normalize = transforms.Normalize(mean=self.input_cfg.mean_val, std=self.input_cfg.std_val)
+            #to_rgb = transforms.Lambda(lambda image: image.convert('RGB'))
+            transform = transforms.Compose([transforms.ToTensor(), normalize,])
+        return transform
+
+    def __getitem__(self, id):
+
+        img_path, label = self.fimgs[id], self.labels[id]
+
+        img = Image.open(f'{self.fdir_cfg.PATH2GTAUAV_IMGS}/{img_path}')
+        img = img.convert('RGB')
+        h_cur, w_cur = img.size
+        #if h_cur < self.h or w_cur < self.w or self._mode != "train":
+        img = img.resize([self.h, self.w])
+
+        if self._mode == "train":
+            img_tensor = self.transform_op(image=np.array(img))["image"]
+            #img_tensor = self.transform_op_train(img_tensor)
+        else:
+            img_tensor = self.transform_op(img)
+
+        lab_tensor = self.to_one_hot([label])
+        #oh_lab = self.to_one_hot(label)
+        return img_tensor, lab_tensor
+
+    def __len__(self):
+        return len(self.labels)
+
+    def to_one_hot(self, values):
+        value_ids = self.lab2id.transform(values)
+        #oh = torch.eye(len(self.lab2id.classes_))[value_ids]
+        #oh = oh.type(torch.int64)
+        #oh = torch.reshape(oh, (-1,))
+        value_ids = torch.as_tensor(value_ids[0], dtype=torch.int64)
+        return value_ids
+
+    def show_image(self, index):
+        """
+        displays the image
+        """
+        image, target = self.__getitem__(index)
+        ts.show(image)
+        ts.save(image, path='1.png')
+
+    def _init_dataset(self):
+        if self._mode == 'train':
+            path2csv = self.fdir_cfg.gtauav_train_csv
+        elif self._mode == 'test':
+            path2csv = self.fdir_cfg.gtauav_test_csv
+
+        self.logger.info(f'Loading csv {self._mode} data ')
+        df = pd.read_csv(path2csv, sep=",", header=0, names=self.fdir_cfg.gtauav_columns)
+        self.fimgs = df['img_name'].values.tolist()
+        self.labels = df['drone_yaw'].values.tolist()
+        self.lab2id.fit(list(self.labels))
+        self.logger.info(f'Done loading {self._mode} data !')
+
+
+def test():
+    from src.conf.base_conf import get_fdir_cfg, get_inputval_cfg
+    path2cfg = fr'{get_proj_dir()}in/config_files/'
+    fdir_conf = get_fdir_cfg(path2cfg)
+    input_conf = get_inputval_cfg(path2cfg)
+
+
+    from torch.utils.data import DataLoader
+
+
+    uav_ds = ObjectUAVDataset(fdir_cfg=fdir_conf, input_cfg=input_conf, mode="test")
+    dl = DataLoader(uav_ds, batch_size=8, shuffle=False, num_workers=6, pin_memory=True)
+    for i in range(10):
+        for i in dl:
+            img_b, lab_b = i
+            print(lab_b)
+            print(img_b.shape)
+
+
+if __name__ == "__main__":
+    test()

src/ds/naruto_dl.py

@@ -0,0 +1,257 @@
+
+
+
+
+
+
+import torch
+
+
+
+# ObjectNet Dataset: Reanalysis and Correction
+# https://github.com/aliborji/ObjectNetReanalysis
+
+
+import json
+
+
+from torchvision.datasets import ImageFolder
+import os
+from PIL import Image
+from PIL import ImageOps
+import torchvision.transforms as transforms
+import torchvision
+from pathlib import Path
+import torch
+from torchvision.utils import draw_bounding_boxes
+
+import pandas as pd
+import numpy as np
+
+from tqdm import tqdm
+import json
+from torch.utils.data import Dataset
+import copy
+import dask.dataframe as dd
+from sklearn.preprocessing import LabelEncoder
+from PIL import Image, ImageFile
+import torchshow as ts
+import matplotlib.pyplot as plt
+import shutil
+from natsort import natsorted
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from src.utils.utils_file_dir import get_files
+
+
+
+####
+
+from src.utils.utils_log import create_logger
+from src.utils.utils_file_dir import create_dir_tree, get_proj_dir
+from src.conf.base_conf import get_fdir_cfg
+
+
+from torchvision.transforms.functional import pil_to_tensor
+from torchvision.io import decode_image
+
+# https://stackoverflow.com/questions/60584155/oserror-image-file-is-truncated
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+torch.manual_seed(0)
+
+
+class ObjectNarutoSign():
+    def __init__(self, fdir_conf):
+
+
+        self.conf = fdir_conf
+
+    def ext_info(self, mode):
+        """
+
+        :param fname:
+        :return:
+        """
+
+        arr_imgs = []
+        arr_labels = []
+        csv_out = None
+
+        if mode == 'train':
+            csv_out = self.conf.naruto_train_csv
+            img_fdir = f'{self.conf.PATH2NARUTO}//train'
+        elif mode == 'test':
+            csv_out = self.conf.naruto_test_csv
+            img_fdir = f'{self.conf.PATH2NARUTO}//test'
+
+
+        for i in self.conf.NARUTO_LABELS:
+            t_imgs = get_files(f'{img_fdir}//{i}', 'jpg')
+            t_labels = [i] * len(t_imgs)
+            arr_imgs.extend(t_imgs)
+            arr_labels.extend(t_labels)
+
+
+        df = pd.DataFrame(data={'img_name': arr_imgs,
+                                'labels': arr_labels})
+        print('Write to dataframe!')
+        df.to_csv(csv_out, sep='\t', columns=['img_name', 'labels'], index=False)
+
+    def main(self):
+        """
+        :return:
+        """
+        arr_mode = ['train', 'test']
+        for mode in arr_mode:
+            self.ext_info(mode)
+
+
+class ObjectNarutoDataset(Dataset):
+    def __init__(self, mode, input_cfg, fdir_cfg):
+        assert mode in ['train', 'test'], 'Parameter mode is wrong!!!'
+        self._mode = mode
+        self.input_cfg = input_cfg
+        self.fdir_cfg = fdir_cfg
+        #
+        self.logger = create_logger(log_name=type(self).__name__)
+        #
+        self.labels = []
+        self.fimgs = []
+        self.lab2id = LabelEncoder()
+        #
+        if self._mode == "train":
+            self.h = input_cfg.height_train
+            self.w = input_cfg.width_train
+            self.transform_op_train = self.transforms_op_pt_train()
+        else:
+            self.h = input_cfg.height_test
+            self.w = input_cfg.width_test
+
+        #
+        self.transform_op = self.init_transforms()
+
+        self._init_dataset()
+
+
+
+    def transforms_op_pt_train(self):
+        '''transform = transforms.Compose([
+            transforms.Resize(224),
+            transforms.RandomErasing(p=0.4),
+            transforms.RandomCrop(224),
+            transforms.RandomInvert(p=0.2),
+            transforms.RandomPosterize(bits=2, p=0.2),
+            transforms.ColorJitter(brightness=.5, hue=.3, contrast=.2, saturation=.2),
+            #transforms.RandomSolarize(p=0.2, threshold=220),
+            transforms.RandomAdjustSharpness(p=0.2, sharpness_factor=2),
+            transforms.RandomAutocontrast(p=0.2),
+            transforms.RandomEqualize(p=0.2),
+            transforms.Normalize(mean=self.input_cfg.mean_val, std=self.input_cfg.std_val),
+            #transforms.RandomPerspective(distortion_scale=0.7, p=0.2, fill=0),
+            #transforms.RandomRotation
+
+        ])'''
+
+        transform = A.Compose([ #A.RandomCrop(self.h, self.w, p=0.2),
+            A.CenterCrop(self.h, self.w, p=0.2),
+            A.Resize(self.h, self.w),
+            #A.RandomRain(p=0.2),
+            #A.RandomSnow(p=0.4),
+            A.RandomFog(p=0.4), #Simulates fog for the image by adding random fog-like artifacts.
+            A.RandomBrightnessContrast(p=0.4),
+            A.RandomGamma(p=0.4),
+            #A.RandomGridShuffle(grid=(16,16), p=0.3),
+            A.RandomShadow(p=0.4),
+            A.RandomSunFlare(p=0.2),
+            A.RandomToneCurve(p=0.2),
+            A.RGBShift(r_shift_limit=25, g_shift_limit=25, b_shift_limit=25, p=0.5),
+            A.Normalize(mean=self.input_cfg.mean_val, std=self.input_cfg.std_val),
+            ToTensorV2(),
+        ])
+        return transform
+
+
+    def init_transforms(self):
+        if self._mode == 'train':
+            transform = self.transforms_op_pt_train()
+        else:
+            normalize = transforms.Normalize(mean=self.input_cfg.mean_val, std=self.input_cfg.std_val)
+            #to_rgb = transforms.Lambda(lambda image: image.convert('RGB'))
+            transform = transforms.Compose([transforms.ToTensor(), normalize,])
+        return transform
+
+    def __getitem__(self, id):
+
+        img_path, label = self.fimgs[id], self.labels[id]
+
+        img = Image.open(img_path)
+        img = img.convert('RGB')
+        h_cur, w_cur = img.size
+        #if h_cur < self.h or w_cur < self.w or self._mode != "train":
+        img = img.resize([self.h, self.w])
+
+        if self._mode == "train":
+            img_tensor = self.transform_op(image=np.array(img))["image"]
+            #img_tensor = self.transform_op_train(img_tensor)
+        else:
+            img_tensor = self.transform_op(img)
+
+        lab_tensor = self.to_one_hot([label])
+        #oh_lab = self.to_one_hot(label)
+        return img_tensor, lab_tensor
+
+    def __len__(self):
+        return len(self.labels)
+
+    def to_one_hot(self, values):
+        value_ids = self.lab2id.transform(values)
+        #oh = torch.eye(len(self.lab2id.classes_))[value_ids]
+        #oh = oh.type(torch.int64)
+        #oh = torch.reshape(oh, (-1,))
+        value_ids = torch.as_tensor(value_ids[0], dtype=torch.int64)
+        return value_ids
+
+    def show_image(self, index):
+        """
+        displays the image
+        """
+        image, target = self.__getitem__(index)
+        ts.show(image)
+        ts.save(image, path='1.png')
+
+    def _init_dataset(self):
+        if self._mode == 'train':
+            path2csv = self.fdir_cfg.naruto_train_csv
+        elif self._mode == 'test':
+            path2csv = self.fdir_cfg.naruto_test_csv
+
+        self.logger.info(f'Loading csv {self._mode} data ')
+        df = pd.read_csv(path2csv, sep="\t", header=0, names=self.fdir_cfg.naruto_columns)
+        self.fimgs = df['img_name'].values.tolist()
+        self.labels = df['labels'].values.tolist()
+        self.lab2id.fit(list(self.labels))
+        self.logger.info(f'Done loading {self._mode} data !')
+
+
+def test():
+    from src.conf.base_conf import get_fdir_cfg, get_inputval_cfg
+    path2cfg = fr'{get_proj_dir()}in/config_files/'
+    fdir_conf = get_fdir_cfg(path2cfg)
+    input_conf = get_inputval_cfg(path2cfg)
+
+
+    from torch.utils.data import DataLoader
+
+    uav_ds = ObjectNarutoDataset(fdir_cfg=fdir_conf, input_cfg=input_conf, mode="test")
+    dl = DataLoader(uav_ds, batch_size=8, shuffle=False, num_workers=6, pin_memory=True)
+    for i in range(10):
+        for i in dl:
+            img_b, lab_b = i
+            print(lab_b)
+            print(img_b.shape)
+
+
+if __name__ == "__main__":
+    test()

src/models/bb/coca/coca_model.py

@@ -0,0 +1,28 @@
+import open_clip
+import torch
+from PIL import Image
+
+
+
+
+def get_coca_model():
+    model, _, transform = open_clip.create_model_and_transforms(
+        model_name="coca_ViT-L-14",
+        pretrained="mscoco_finetuned_laion2B-s13B-b90k"
+    )
+    model = model.to('cuda')
+    return model#, transform
+
+
+def test():
+    model = get_coca_model()
+    #im = Image.open("/home/servml/Документы/datasets/objectdet_crop/images/test/air_freshener/1dd007f993b941c.png").convert("RGB")
+    #im = transform(im).unsqueeze(0)
+    im = torch.rand((4, 3 , 224, 224))
+    print(im.shape)
+
+    res = model.encode_image(im)
+    print(res)
+
+if __name__ == "__main__":
+    test()

src/models/bb/edgenext/edgenext_blocks.py

@@ -0,0 +1,335 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+    def extra_repr(self):
+        return f'drop_prob={round(self.drop_prob,3):0.3f}'
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+class PositionalEncodingFourier(nn.Module):
+    def __init__(self, hidden_dim=32, dim=768, temperature=10000):
+        super().__init__()
+        self.token_projection = nn.Conv2d(hidden_dim * 2, dim, kernel_size=1)
+        self.scale = 2 * math.pi
+        self.temperature = temperature
+        self.hidden_dim = hidden_dim
+        self.dim = dim
+
+    def forward(self, B, H, W):
+        mask = torch.zeros(B, H, W).bool().to(self.token_projection.weight.device)
+        not_mask = ~mask
+        y_embed = not_mask.cumsum(1, dtype=torch.float32)
+        x_embed = not_mask.cumsum(2, dtype=torch.float32)
+        eps = 1e-6
+        y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+        x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(self.hidden_dim, dtype=torch.float32, device=mask.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.hidden_dim)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(),
+                             pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(),
+                             pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+        pos = self.token_projection(pos)
+
+        return pos
+
+
+class ConvEncoder(nn.Module):
+    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6, expan_ratio=4, kernel_size=7):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim)
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, expan_ratio * dim)
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(expan_ratio * dim, dim)
+        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(dim),
+                                  requires_grad=True) if layer_scale_init_value > 0 else None
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+        return x
+
+
+class ConvEncoderBNHS(nn.Module):
+    """
+        Conv. Encoder with Batch Norm and Hard-Swish Activation
+    """
+    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6, expan_ratio=4, kernel_size=7):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim, bias=False)
+        self.norm = nn.BatchNorm2d(dim)
+        self.pwconv1 = nn.Linear(dim, expan_ratio * dim)
+        self.act = nn.Hardswish()
+        self.pwconv2 = nn.Linear(expan_ratio * dim, dim)
+        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(dim),
+                                  requires_grad=True) if layer_scale_init_value > 0 else None
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = self.norm(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+        return x
+
+
+class SDTAEncoder(nn.Module):
+    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6, expan_ratio=4,
+                 use_pos_emb=True, num_heads=8, qkv_bias=True, attn_drop=0., drop=0., scales=1):
+        super().__init__()
+        width = max(int(math.ceil(dim / scales)), int(math.floor(dim // scales)))
+        self.width = width
+        if scales == 1:
+            self.nums = 1
+        else:
+            self.nums = scales - 1
+        convs = []
+        for i in range(self.nums):
+            convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, groups=width))
+        self.convs = nn.ModuleList(convs)
+
+        self.pos_embd = None
+        if use_pos_emb:
+            self.pos_embd = PositionalEncodingFourier(dim=dim)
+        self.norm_xca = LayerNorm(dim, eps=1e-6)
+        self.gamma_xca = nn.Parameter(layer_scale_init_value * torch.ones(dim),
+                                      requires_grad=True) if layer_scale_init_value > 0 else None
+        self.xca = XCA(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, expan_ratio * dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()  # TODO: MobileViT is using 'swish'
+        self.pwconv2 = nn.Linear(expan_ratio * dim, dim)
+        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)),
+                                  requires_grad=True) if layer_scale_init_value > 0 else None
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        input = x
+
+        spx = torch.split(x, self.width, 1)
+        for i in range(self.nums):
+            if i == 0:
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            sp = self.convs[i](sp)
+            if i == 0:
+                out = sp
+            else:
+                out = torch.cat((out, sp), 1)
+        x = torch.cat((out, spx[self.nums]), 1)
+        # XCA
+        B, C, H, W = x.shape
+        x = x.reshape(B, C, H * W).permute(0, 2, 1)
+        if self.pos_embd:
+            pos_encoding = self.pos_embd(B, H, W).reshape(B, -1, x.shape[1]).permute(0, 2, 1)
+            x = x + pos_encoding
+        x = x + self.drop_path(self.gamma_xca * self.xca(self.norm_xca(x)))
+        x = x.reshape(B, H, W, C)
+
+        # Inverted Bottleneck
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+
+        return x
+
+
+class SDTAEncoderBNHS(nn.Module):
+    """
+        SDTA Encoder with Batch Norm and Hard-Swish Activation
+    """
+    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6, expan_ratio=4,
+                 use_pos_emb=True, num_heads=8, qkv_bias=True, attn_drop=0., drop=0., scales=1):
+        super().__init__()
+        width = max(int(math.ceil(dim / scales)), int(math.floor(dim // scales)))
+        self.width = width
+        if scales == 1:
+            self.nums = 1
+        else:
+            self.nums = scales - 1
+        convs = []
+        for i in range(self.nums):
+            convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, groups=width))
+        self.convs = nn.ModuleList(convs)
+
+        self.pos_embd = None
+        if use_pos_emb:
+            self.pos_embd = PositionalEncodingFourier(dim=dim)
+        self.norm_xca = nn.BatchNorm2d(dim)
+        self.gamma_xca = nn.Parameter(layer_scale_init_value * torch.ones(dim),
+                                      requires_grad=True) if layer_scale_init_value > 0 else None
+        self.xca = XCA(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+
+        self.norm = nn.BatchNorm2d(dim)
+        self.pwconv1 = nn.Linear(dim, expan_ratio * dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.Hardswish()  # TODO: MobileViT is using 'swish'
+        self.pwconv2 = nn.Linear(expan_ratio * dim, dim)
+        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)),
+                                  requires_grad=True) if layer_scale_init_value > 0 else None
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        input = x
+
+        spx = torch.split(x, self.width, 1)
+        for i in range(self.nums):
+            if i == 0:
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            sp = self.convs[i](sp)
+            if i == 0:
+                out = sp
+            else:
+                out = torch.cat((out, sp), 1)
+        x = torch.cat((out, spx[self.nums]), 1)
+        # XCA
+        x = self.norm_xca(x)
+        B, C, H, W = x.shape
+        x = x.reshape(B, C, H * W).permute(0, 2, 1)
+        if self.pos_embd:
+            pos_encoding = self.pos_embd(B, H, W).reshape(B, -1, x.shape[1]).permute(0, 2, 1)
+            x = x + pos_encoding
+        x = x + self.drop_path(self.gamma_xca * self.xca(x))
+        x = x.reshape(B, H, W, C).permute(0, 3, 1, 2)
+
+        # Inverted Bottleneck
+        x = self.norm(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+
+        return x
+
+
+class XCA(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q.transpose(-2, -1)
+        k = k.transpose(-2, -1)
+        v = v.transpose(-2, -1)
+
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature
+        # -------------------
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).permute(0, 3, 1, 2).reshape(B, N, C)
+        # ------------------
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'temperature'}

src/models/bb/edgenext/edgenext_model.py

@@ -0,0 +1,401 @@
+# https://github.com/mmaaz60/EdgeNeXt
+# https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/edgenext.py
+
+from src.models.bb.edgenext.edgenext_blocks import *
+import torch
+from torch import nn
+from timm.models.layers import trunc_normal_
+
+
+
+
+class EdgeNeXtBNHS(nn.Module):
+    def __init__(self, in_chans=3, num_classes=1000,
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+                 global_block=[0, 0, 0, 3], global_block_type=['None', 'None', 'None', 'SDTA_BN_HS'],
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., expan_ratio=4,
+                 kernel_sizes=[7, 7, 7, 7], heads=[8, 8, 8, 8], use_pos_embd_xca=[False, False, False, False],
+                 use_pos_embd_global=False, d2_scales=[2, 3, 4, 5], **kwargs):
+        super().__init__()
+        for g in global_block_type:
+            assert g in ['None', 'SDTA_BN_HS']
+
+        if use_pos_embd_global:
+            self.pos_embd = PositionalEncodingFourier(dim=dims[0])
+        else:
+            self.pos_embd = None
+
+        self.downsample_layers = nn.ModuleList()  # stem and 3 intermediate downsampling conv layers
+        stem = nn.Sequential(
+            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4, bias=False),
+            nn.BatchNorm2d(dims[0])
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = nn.Sequential(
+                nn.BatchNorm2d(dims[i]),
+                nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2, bias=False),
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = nn.ModuleList()  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(4):
+            stage_blocks = []
+            for j in range(depths[i]):
+                if j > depths[i] - global_block[i] - 1:
+                    if global_block_type[i] == 'SDTA_BN_HS':
+                        stage_blocks.append(SDTAEncoderBNHS(dim=dims[i], drop_path=dp_rates[cur + j],
+                                                            expan_ratio=expan_ratio, scales=d2_scales[i],
+                                                            use_pos_emb=use_pos_embd_xca[i],
+                                                            num_heads=heads[i], ))
+                    else:
+                        raise NotImplementedError
+                else:
+                    stage_blocks.append(ConvEncoderBNHS(dim=dims[i], drop_path=dp_rates[cur + j],
+                                                        layer_scale_init_value=layer_scale_init_value,
+                                                        expan_ratio=expan_ratio, kernel_size=kernel_sizes[i]))
+
+            self.stages.append(nn.Sequential(*stage_blocks))
+            cur += depths[i]
+        self.norm = nn.BatchNorm2d(dims[-1])
+        self.head = nn.Linear(dims[-1], num_classes)
+
+        self.apply(self._init_weights)
+        self.head_dropout = nn.Dropout(kwargs["classifier_dropout"])
+        self.head.weight.data.mul_(head_init_scale)
+        self.head.bias.data.mul_(head_init_scale)
+
+    def _init_weights(self, m):  # TODO: MobileViT is using 'kaiming_normal' for initializing conv layers
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (LayerNorm, nn.LayerNorm)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward_features(self, x):
+        x = self.downsample_layers[0](x)
+        x = self.stages[0](x)
+        if self.pos_embd:
+            B, C, H, W = x.shape
+            x = x + self.pos_embd(B, H, W)
+        for i in range(1, 4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+        return self.norm(x).mean([-2, -1])
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(self.head_dropout(x))
+        return
+
+
+class EdgeNeXt(nn.Module):
+    def __init__(self, in_chans=3, num_classes=1000,
+                 depths=[3, 3, 9, 3], dims=[24, 48, 88, 168],
+                 global_block=[0, 0, 0, 3], global_block_type=['None', 'None', 'None', 'SDTA'],
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., expan_ratio=4,
+                 kernel_sizes=[7, 7, 7, 7], heads=[8, 8, 8, 8], use_pos_embd_xca=[False, False, False, False],
+                 use_pos_embd_global=False, d2_scales=[2, 3, 4, 5], class_drop_rate=0.2, **kwargs):
+        super().__init__()
+        for g in global_block_type:
+            assert g in ['None', 'SDTA']
+        if use_pos_embd_global:
+            self.pos_embd = PositionalEncodingFourier(dim=dims[0])
+        else:
+            self.pos_embd = None
+        self.downsample_layers = nn.ModuleList()  # stem and 3 intermediate downsampling conv layers
+        stem = nn.Sequential(
+            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+            LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = nn.Sequential(
+                LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
+                nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = nn.ModuleList()  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(4):
+            stage_blocks = []
+            for j in range(depths[i]):
+                if j > depths[i] - global_block[i] - 1:
+                    if global_block_type[i] == 'SDTA':
+                        stage_blocks.append(SDTAEncoder(dim=dims[i], drop_path=dp_rates[cur + j],
+                                                        expan_ratio=expan_ratio, scales=d2_scales[i],
+                                                        use_pos_emb=use_pos_embd_xca[i], num_heads=heads[i]))
+                    else:
+                        raise NotImplementedError
+                else:
+                    stage_blocks.append(ConvEncoder(dim=dims[i], drop_path=dp_rates[cur + j],
+                                                    layer_scale_init_value=layer_scale_init_value,
+                                                    expan_ratio=expan_ratio, kernel_size=kernel_sizes[i]))
+
+            self.stages.append(nn.Sequential(*stage_blocks))
+            cur += depths[i]
+        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)  # Final norm layer
+        self.head = nn.Linear(dims[-1], num_classes)
+
+        self.apply(self._init_weights)
+        self.head_dropout = nn.Dropout(class_drop_rate)
+        self.head.weight.data.mul_(head_init_scale)
+        self.head.bias.data.mul_(head_init_scale)
+
+    def _init_weights(self, m):  # TODO: MobileViT is using 'kaiming_normal' for initializing conv layers
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (LayerNorm, nn.LayerNorm)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward_features(self, x):
+        x = self.downsample_layers[0](x)
+        x = self.stages[0](x)
+        if self.pos_embd:
+            B, C, H, W = x.shape
+            x = x + self.pos_embd(B, H, W)
+        for i in range(1, 4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+
+        return self.norm(x.mean([-2, -1]))  # Global average pooling, (N, C, H, W) -> (N, C)
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(self.head_dropout(x))
+        return x
+
+
+class EdgeNeXtBNHS(nn.Module):
+    def __init__(self, in_chans=3, num_classes=1000,
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+                 global_block=[0, 0, 0, 3], global_block_type=['None', 'None', 'None', 'SDTA_BN_HS'],
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., expan_ratio=4,
+                 kernel_sizes=[7, 7, 7, 7], heads=[8, 8, 8, 8], use_pos_embd_xca=[False, False, False, False],
+                 use_pos_embd_global=False, d2_scales=[2, 3, 4, 5], class_drop_rate=0.2, **kwargs):
+        super().__init__()
+        for g in global_block_type:
+            assert g in ['None', 'SDTA_BN_HS']
+
+        if use_pos_embd_global:
+            self.pos_embd = PositionalEncodingFourier(dim=dims[0])
+        else:
+            self.pos_embd = None
+
+        self.downsample_layers = nn.ModuleList()  # stem and 3 intermediate downsampling conv layers
+        stem = nn.Sequential(
+            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4, bias=False),
+            nn.BatchNorm2d(dims[0])
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = nn.Sequential(
+                nn.BatchNorm2d(dims[i]),
+                nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2, bias=False),
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = nn.ModuleList()  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(4):
+            stage_blocks = []
+            for j in range(depths[i]):
+                if j > depths[i] - global_block[i] - 1:
+                    if global_block_type[i] == 'SDTA_BN_HS':
+                        stage_blocks.append(SDTAEncoderBNHS(dim=dims[i], drop_path=dp_rates[cur + j],
+                                                            expan_ratio=expan_ratio, scales=d2_scales[i],
+                                                            use_pos_emb=use_pos_embd_xca[i],
+                                                            num_heads=heads[i]))
+                    else:
+                        raise NotImplementedError
+                else:
+                    stage_blocks.append(ConvEncoderBNHS(dim=dims[i], drop_path=dp_rates[cur + j],
+                                                        layer_scale_init_value=layer_scale_init_value,
+                                                        expan_ratio=expan_ratio, kernel_size=kernel_sizes[i]))
+
+            self.stages.append(nn.Sequential(*stage_blocks))
+            cur += depths[i]
+        self.norm = nn.BatchNorm2d(dims[-1])
+        self.head = nn.Linear(dims[-1], num_classes)
+
+        self.apply(self._init_weights)
+        self.head_dropout = nn.Dropout(class_drop_rate)
+        self.head.weight.data.mul_(head_init_scale)
+        self.head.bias.data.mul_(head_init_scale)
+
+    def _init_weights(self, m):  # TODO: MobileViT is using 'kaiming_normal' for initializing conv layers
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (LayerNorm, nn.LayerNorm)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward_features(self, x):
+        x = self.downsample_layers[0](x)
+        x = self.stages[0](x)
+        if self.pos_embd:
+            B, C, H, W = x.shape
+            x = x + self.pos_embd(B, H, W)
+        for i in range(1, 4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+        return self.norm(x).mean([-2, -1])
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(self.head_dropout(x))
+        return x
+
+
+
+
+
+
+"""
+-- Main Models
+    XX-Small -> 1.3M
+    X-Small -> 2.3M
+    Small -> 5.6M
+"""
+
+
+
+
+
+
+def edgenext_xx_small(pretrained=False, **kwargs):
+    # 1.33M & 260.58M @ 256 resolution
+    # 71.23% Top-1 accuracy
+    # No AA, Color Jitter=0.4, No Mixup & Cutmix, DropPath=0.0, BS=4096, lr=0.006, multi-scale-sampler
+    # Jetson FPS=51.66 versus 47.67 for MobileViT_XXS
+    # For A100: FPS @ BS=1: 212.13 & @ BS=256: 7042.06 versus FPS @ BS=1: 96.68 & @ BS=256: 4624.71 for MobileViT_XXS
+    model = EdgeNeXt(depths=[2, 2, 6, 2], dims=[24, 48, 88, 168], expan_ratio=4,
+                     global_block=[0, 1, 1, 1],
+                     global_block_type=['None', 'SDTA', 'SDTA', 'SDTA'],
+                     use_pos_embd_xca=[False, True, False, False],
+                     kernel_sizes=[3, 5, 7, 9],
+                     heads=[4, 4, 4, 4],
+                     d2_scales=[2, 2, 3, 4],  drop_path_rate=0.0,
+                     **kwargs)
+
+    return model
+
+
+
+def edgenext_x_small(pretrained=False, **kwargs):
+    # 2.34M & 538.0M @ 256 resolution
+    # 75.00% Top-1 accuracy
+    # No AA, No Mixup & Cutmix, DropPath=0.0, BS=4096, lr=0.006, multi-scale-sampler
+    # Jetson FPS=31.61 versus 28.49 for MobileViT_XS
+    # For A100: FPS @ BS=1: 179.55 & @ BS=256: 4404.95 versus FPS @ BS=1: 94.55 & @ BS=256: 2361.53 for MobileViT_XS
+    model = EdgeNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 100, 192], expan_ratio=4,
+                     global_block=[0, 1, 1, 1],
+                     global_block_type=['None', 'SDTA', 'SDTA', 'SDTA'],
+                     use_pos_embd_xca=[False, True, False, False],
+                     kernel_sizes=[3, 5, 7, 9],
+                     heads=[4, 4, 4, 4],
+                     d2_scales=[2, 2, 3, 4],
+                     **kwargs)
+
+    return model
+
+
+
+def edgenext_small(pretrained=False, **kwargs):
+    # 5.59M & 1260.59M @ 256 resolution
+    # 79.43% Top-1 accuracy
+    # AA=True, No Mixup & Cutmix, DropPath=0.1, BS=4096, lr=0.006, multi-scale-sampler
+    # Jetson FPS=20.47 versus 18.86 for MobileViT_S
+    # For A100: FPS @ BS=1: 172.33 & @ BS=256: 3010.25 versus FPS @ BS=1: 93.84 & @ BS=256: 1785.92 for MobileViT_S
+    model = EdgeNeXt(depths=[3, 3, 9, 3], dims=[48, 96, 160, 304], expan_ratio=4,
+                     global_block=[0, 1, 1, 1],
+                     global_block_type=['None', 'SDTA', 'SDTA', 'SDTA'],
+                     use_pos_embd_xca=[False, True, False, False],
+                     kernel_sizes=[3, 5, 7, 9],
+                     d2_scales=[2, 2, 3, 4],
+                     **kwargs)
+
+    return model
+
+
+
+def edgenext_base(pretrained=False, **kwargs):
+    # 18.51M & 3840.93M @ 256 resolution
+    # 82.5% (normal) 83.7% (USI) Top-1 accuracy
+    # AA=True, Mixup & Cutmix, DropPath=0.1, BS=4096, lr=0.006, multi-scale-sampler
+    # Jetson FPS=xx.xx versus xx.xx for MobileViT_S
+    # For A100: FPS @ BS=1: xxx.xx & @ BS=256: xxxx.xx
+    model = EdgeNeXt(depths=[3, 3, 9, 3], dims=[80, 160, 288, 584], expan_ratio=4,
+                     global_block=[0, 1, 1, 1],
+                     global_block_type=['None', 'SDTA', 'SDTA', 'SDTA'],
+                     use_pos_embd_xca=[False, True, False, False],
+                     kernel_sizes=[3, 5, 7, 9],
+                     d2_scales=[2, 2, 3, 4], drop_path_rate=0.1,
+                     **kwargs)
+
+    return model
+
+"""
+    Using BN & HSwish instead of LN & GeLU
+"""
+
+
+def edgenext_xx_small_bn_hs(pretrained=False, **kwargs):
+    # 1.33M & 259.53M @ 256 resolution
+    # 70.33% Top-1 accuracy
+    # For A100: FPS @ BS=1: 219.66 & @ BS=256: 10359.98
+    model = EdgeNeXtBNHS(depths=[2, 2, 6, 2], dims=[24, 48, 88, 168], expan_ratio=4,
+                         global_block=[0, 1, 1, 1],
+                         global_block_type=['None', 'SDTA_BN_HS', 'SDTA_BN_HS', 'SDTA_BN_HS'],
+                         use_pos_embd_xca=[False, True, False, False],
+                         kernel_sizes=[3, 5, 7, 9],
+                         heads=[4, 4, 4, 4],
+                         d2_scales=[2, 2, 3, 4],
+                         **kwargs)
+
+    return model
+
+
+
+def edgenext_x_small_bn_hs(pretrained=False, **kwargs):
+    # 2.34M & 535.84M @ 256 resolution
+    # 74.87% Top-1 accuracy
+    # For A100: FPS @ BS=1: 179.25 & @ BS=256: 6059.59
+    model = EdgeNeXtBNHS(depths=[3, 3, 9, 3], dims=[32, 64, 100, 192], expan_ratio=4,
+                         global_block=[0, 1, 1, 1],
+                         global_block_type=['None', 'SDTA_BN_HS', 'SDTA_BN_HS', 'SDTA_BN_HS'],
+                         use_pos_embd_xca=[False, True, False, False],
+                         kernel_sizes=[3, 5, 7, 9],
+                         heads=[4, 4, 4, 4],
+                         d2_scales=[2, 2, 3, 4], drop_path_rate=0.2,
+                         **kwargs)
+
+    return model
+
+
+
+def edgenext_small_bn_hs(pretrained=False, **kwargs):
+    # 5.58M & 1257.28M @ 256 resolution
+    # 78.39% Top-1 accuracy
+    # For A100: FPS @ BS=1: 174.68 & @ BS=256: 3808.19
+    model = EdgeNeXtBNHS(depths=[3, 3, 9, 3], dims=[48, 96, 160, 304], expan_ratio=4,
+                         global_block=[0, 1, 1, 1],
+                         global_block_type=['None', 'SDTA_BN_HS', 'SDTA_BN_HS', 'SDTA_BN_HS'],
+                         use_pos_embd_xca=[False, True, False, False],
+                         kernel_sizes=[3, 5, 7, 9],
+                         d2_scales=[2, 2, 3, 4], drop_path_rate=0.2,
+                         **kwargs)
+
+    return model

src/models/bb/inception_next/inception_next_model.py

@@ -0,0 +1,323 @@
+"""
+InceptionNeXt implementation, paper: https://arxiv.org/abs/2303.16900
+
+Some code is borrowed from timm: https://github.com/huggingface/pytorch-image-models
+"""
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.models.helpers import checkpoint_seq
+from timm.models.layers import trunc_normal_, DropPath
+from timm.models.registry import register_model
+from timm.layers.helpers import to_2tuple
+
+
+class InceptionDWConv2d(nn.Module):
+    """ Inception depthweise convolution
+    """
+    def __init__(self, in_channels, square_kernel_size=3, band_kernel_size=11, branch_ratio=0.125):
+        super().__init__()
+
+        gc = int(in_channels * branch_ratio) # channel numbers of a convolution branch
+        self.dwconv_hw = nn.Conv2d(gc, gc, square_kernel_size, padding=square_kernel_size//2, groups=gc)
+        self.dwconv_w = nn.Conv2d(gc, gc, kernel_size=(1, band_kernel_size), padding=(0, band_kernel_size//2), groups=gc)
+        self.dwconv_h = nn.Conv2d(gc, gc, kernel_size=(band_kernel_size, 1), padding=(band_kernel_size//2, 0), groups=gc)
+        self.split_indexes = (in_channels - 3 * gc, gc, gc, gc)
+
+    def forward(self, x):
+        x_id, x_hw, x_w, x_h = torch.split(x, self.split_indexes, dim=1)
+        return torch.cat(
+            (x_id, self.dwconv_hw(x_hw), self.dwconv_w(x_w), self.dwconv_h(x_h)),
+            dim=1,
+        )
+
+
+class ConvMlp(nn.Module):
+    """ MLP using 1x1 convs that keeps spatial dims
+    copied from timm: https://github.com/huggingface/pytorch-image-models/blob/v0.6.11/timm/models/layers/mlp.py
+    """
+    def __init__(
+            self, in_features, hidden_features=None, out_features=None, act_layer=nn.ReLU,
+            norm_layer=None, bias=True, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        bias = to_2tuple(bias)
+
+        self.fc1 = nn.Conv2d(in_features, hidden_features, kernel_size=1, bias=bias[0])
+        self.norm = norm_layer(hidden_features) if norm_layer else nn.Identity()
+        self.act = act_layer()
+        self.drop = nn.Dropout(drop)
+        self.fc2 = nn.Conv2d(hidden_features, out_features, kernel_size=1, bias=bias[1])
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.norm(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        return x
+
+
+class MlpHead(nn.Module):
+    """ MLP classification head
+    """
+    def __init__(self, dim, num_classes=1000, mlp_ratio=3, act_layer=nn.GELU,
+                 norm_layer=partial(nn.LayerNorm, eps=1e-6), drop=0., bias=True):
+        super().__init__()
+        hidden_features = int(mlp_ratio * dim)
+        self.fc1 = nn.Linear(dim, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.norm = norm_layer(hidden_features)
+        self.fc2 = nn.Linear(hidden_features, num_classes, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = x.mean((2, 3)) # global average pooling
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.norm(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        return x
+
+
+class MetaNeXtBlock(nn.Module):
+    """ MetaNeXtBlock Block
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+        ls_init_value (float): Init value for Layer Scale. Default: 1e-6.
+    """
+
+    def __init__(
+            self,
+            dim,
+            token_mixer=nn.Identity,
+            norm_layer=nn.BatchNorm2d,
+            mlp_layer=ConvMlp,
+            mlp_ratio=4,
+            act_layer=nn.GELU,
+            ls_init_value=1e-6,
+            drop_path=0.,
+
+    ):
+        super().__init__()
+        self.token_mixer = token_mixer(dim)
+        self.norm = norm_layer(dim)
+        self.mlp = mlp_layer(dim, int(mlp_ratio * dim), act_layer=act_layer)
+        self.gamma = nn.Parameter(ls_init_value * torch.ones(dim)) if ls_init_value else None
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        shortcut = x
+        x = self.token_mixer(x)
+        x = self.norm(x)
+        x = self.mlp(x)
+        if self.gamma is not None:
+            x = x.mul(self.gamma.reshape(1, -1, 1, 1))
+        x = self.drop_path(x) + shortcut
+        return x
+
+
+class MetaNeXtStage(nn.Module):
+    def __init__(
+            self,
+            in_chs,
+            out_chs,
+            ds_stride=2,
+            depth=2,
+            drop_path_rates=None,
+            ls_init_value=1.0,
+            token_mixer=nn.Identity,
+            act_layer=nn.GELU,
+            norm_layer=None,
+            mlp_ratio=4,
+    ):
+        super().__init__()
+        self.grad_checkpointing = False
+        if ds_stride > 1:
+            self.downsample = nn.Sequential(
+                norm_layer(in_chs),
+                nn.Conv2d(in_chs, out_chs, kernel_size=ds_stride, stride=ds_stride),
+            )
+        else:
+            self.downsample = nn.Identity()
+
+        drop_path_rates = drop_path_rates or [0.] * depth
+        stage_blocks = []
+        for i in range(depth):
+            stage_blocks.append(MetaNeXtBlock(
+                dim=out_chs,
+                drop_path=drop_path_rates[i],
+                ls_init_value=ls_init_value,
+                token_mixer=token_mixer,
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                mlp_ratio=mlp_ratio,
+            ))
+            in_chs = out_chs
+        self.blocks = nn.Sequential(*stage_blocks)
+
+    def forward(self, x):
+        x = self.downsample(x)
+        if self.grad_checkpointing and not torch.jit.is_scripting():
+            x = checkpoint_seq(self.blocks, x)
+        else:
+            x = self.blocks(x)
+        return x
+
+
+class MetaNeXt(nn.Module):
+    r""" MetaNeXt
+        A PyTorch impl of : `InceptionNeXt: When Inception Meets ConvNeXt`  - https://arxiv.org/pdf/2203.xxxxx.pdf
+
+    Args:
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: (3, 3, 9, 3)
+        dims (tuple(int)): Feature dimension at each stage. Default: (96, 192, 384, 768)
+        token_mixers: Token mixer function. Default: nn.Identity
+        norm_layer: Normalziation layer. Default: nn.BatchNorm2d
+        act_layer: Activation function for MLP. Default: nn.GELU
+        mlp_ratios (int or tuple(int)): MLP ratios. Default: (4, 4, 4, 3)
+        head_fn: classifier head
+        drop_rate (float): Head dropout rate
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        ls_init_value (float): Init value for Layer Scale. Default: 1e-6.
+    """
+
+    def __init__(
+            self,
+            in_chans=3,
+            num_classes=1000,
+            depths=(3, 3, 9, 3),
+            dims=(96, 192, 384, 768),
+            token_mixers=nn.Identity,
+            norm_layer=nn.BatchNorm2d,
+            act_layer=nn.GELU,
+            mlp_ratios=(4, 4, 4, 3),
+            head_fn=MlpHead,
+            drop_rate=0.,
+            drop_path_rate=0.,
+            ls_init_value=1e-6,
+            **kwargs,
+    ):
+        super().__init__()
+
+        num_stage = len(depths)
+        if not isinstance(token_mixers, (list, tuple)):
+            token_mixers = [token_mixers] * num_stage
+        if not isinstance(mlp_ratios, (list, tuple)):
+            mlp_ratios = [mlp_ratios] * num_stage
+
+
+        self.num_classes = num_classes
+        self.drop_rate = drop_rate
+        self.stem = nn.Sequential(
+            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+            norm_layer(dims[0])
+        )
+
+        self.stages = nn.Sequential()
+        dp_rates = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(depths)).split(depths)]
+        stages = []
+        prev_chs = dims[0]
+        # feature resolution stages, each consisting of multiple residual blocks
+        for i in range(num_stage):
+            out_chs = dims[i]
+            stages.append(MetaNeXtStage(
+                prev_chs,
+                out_chs,
+                ds_stride=2 if i > 0 else 1,
+                depth=depths[i],
+                drop_path_rates=dp_rates[i],
+                ls_init_value=ls_init_value,
+                act_layer=act_layer,
+                token_mixer=token_mixers[i],
+                norm_layer=norm_layer,
+                mlp_ratio=mlp_ratios[i],
+            ))
+            prev_chs = out_chs
+        self.stages = nn.Sequential(*stages)
+        self.num_features = prev_chs
+        self.head = head_fn(self.num_features, num_classes, drop=drop_rate)
+        self.apply(self._init_weights)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        for s in self.stages:
+            s.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'norm'}
+
+
+    def forward_features(self, x):
+        x = self.stem(x)
+        x = self.stages(x)
+        return x
+
+    def forward_head(self, x):
+        x = self.head(x)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.forward_head(x)
+        return x
+
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+def _cfg():
+    return {
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7),
+        'crop_pct': 0.875, 'interpolation': 'bicubic',
+        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
+        'first_conv': 'stem.0', 'classifier': 'head.fc',
+    }
+
+def inceptionnext_tiny():
+    model = MetaNeXt(depths=(3, 3, 9, 3), dims=(96, 192, 384, 768),
+                     token_mixers=InceptionDWConv2d,
+                     )
+    return model
+
+
+
+
+
+def inceptionnext_base_384():
+    model = MetaNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024],
+                     mlp_ratios=[4, 4, 4, 3],
+                     token_mixers=InceptionDWConv2d,
+                     )
+    return model
+
+
+def inceptionnext_base():
+    model = MetaNeXt(depths=(3, 3, 27, 3), dims=(128, 256, 512, 1024),
+                     token_mixers=InceptionDWConv2d,
+                     )
+    return model
+
+def get_inception_next_model(model_name):
+    if model_name == 'inception_next_tiny':
+        return inceptionnext_tiny()
+    elif model_name == 'inception_next_base_384':
+        return inceptionnext_base_384()
+    elif model_name == 'inception_next_base':
+        return inceptionnext_base()
+
+if __name__ == "__main__":
+    pass

src/models/bb/load_models.py

@@ -0,0 +1,88 @@
+"""
+Load pretrained models
+"""
+
+from torchvision import models
+import torch
+import torch.nn as nn
+import os
+from src.utils.utils_file_dir import get_proj_dir
+from src.conf.nn_conf import get_nn_cfg
+
+
+from src.models.bb.inception_next.inception_next_model import get_inception_next_model
+
+from src.models.bb.coca.coca_model import get_coca_model
+
+from src.models.bb.edgenext.edgenext_model import edgenext_xx_small_bn_hs, edgenext_base
+
+def load_weights_from_ckpt(nn_conf, model_ft ):
+    """
+    :param model_name:
+    :param model_ft:
+    :param device:
+    :return:
+    """
+
+    device=torch.device(nn_conf.device)
+    proj_dir = get_proj_dir()
+    path2ckpt_model = f'{proj_dir}{nn_conf.path2ckpt}{nn_conf.model_name}.pth'
+
+    ckpt = torch.load(path2ckpt_model, map_location=device, weights_only=False)
+
+    state_dict = None
+    if type(ckpt) == dict:
+        if 'model' in ckpt.keys():
+            state_dict = ckpt['model']
+        elif 'state_dict' in ckpt.keys():
+            state_dict = ckpt['state_dict']
+        else:
+            state_dict = ckpt
+    else:
+        state_dict = ckpt
+
+    try:
+        model_ft.load_state_dict(state_dict)
+    except:
+        missing_keys, unexpected_keys = model_ft.load_state_dict(torch.load(path2ckpt_model), False)
+    return model_ft
+
+def load_models(nn_conf):
+    """
+    :param model_name:
+    :return:
+    """
+    if nn_conf.model_name == 'inception_next_tiny' or nn_conf.model_name == 'inception_next_base_384' or nn_conf.model_name == 'inception_next_base':
+        model_ft = get_inception_next_model(nn_conf.model_name)
+        model_ft = load_weights_from_ckpt(nn_conf, model_ft)
+    elif nn_conf.model_name == 'edgenext_xx_small_bn_hs':
+        model_ft = edgenext_xx_small_bn_hs()
+        model_ft = load_weights_from_ckpt(nn_conf, model_ft)
+    elif nn_conf.model_name == 'edgenext_small_usi':
+        model_ft = edgenext_base()
+        model_ft = load_weights_from_ckpt(nn_conf, model_ft)
+    elif nn_conf.model_name == 'edgenext_base_usi':
+        model_ft = edgenext_base()
+        model_ft = load_weights_from_ckpt(nn_conf, model_ft)
+    elif nn_conf.model_name == 'coca':
+        model_ft = get_coca_model()
+    return model_ft
+
+def test1():
+    path2cfg = fr'{get_proj_dir()}in/config_files/'
+    nn_conf = get_nn_cfg(path2cfg)
+    device = nn_conf.device
+    model = load_models(nn_conf).to(device)
+    print(model)
+    tensor = torch.rand((2, 3, 224, 224)).to(device)
+    res = model(tensor)
+    print(res)
+    del model
+
+if __name__ == "__main__":
+    test1()
+
+
+
+
+

src/models/bb/make_fe.py

@@ -0,0 +1,136 @@
+'''
+format nn from classification task of ImageNet to feature extractors for new task
+'''
+
+import torch
+import torch.nn as nn
+from src.models.bb.load_models import load_models
+from src.utils.utils_file_dir import get_proj_dir
+from src.conf.nn_conf import get_nn_cfg
+from functools import partial
+
+from src.utils.utils_nn import autopad
+
+class MLP(nn.Module):
+    """ MLP classification head
+    """
+    def __init__(self, dim, num_classes=1000, mlp_ratio=3, act_layer=nn.GELU,
+                 norm_layer=partial(nn.LayerNorm, eps=1e-6), drop=0.2, bias=True, use_pooling=False,
+                 use_ext=False, model_ext=None):
+        super().__init__()
+        self.use_ext = use_ext
+        self.model_ext = model_ext
+        self.use_pooling = use_pooling
+        hidden_features = int(mlp_ratio * dim)
+        self.fc1 = nn.Linear(dim, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.norm = norm_layer(hidden_features)
+        self.fc2 = nn.Linear(hidden_features, num_classes, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        #print(x.shape)
+        if self.use_pooling:
+            x = x.mean((2, 3)) # global average pooling
+        if self.use_ext:
+            #x = x.unsqueeze(0)
+            #x = x.transpose((2, 0, 1))
+            #x = torch.randn(32, 3, 224, 224).to('cuda')
+            x = self.model_ext.encode_image(x)
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.norm(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        return x
+
+
+def set_parameter_requires_grad_old(model):
+    #if feature_extracting:
+    for name, param in model.named_parameters():
+        if name.find('head'):
+            param.requires_grad = True
+        else:
+            param.requires_grad = False
+
+
+def set_parameter_requires_grad(model):
+    #if feature_extracting:
+    for name, param in model.named_parameters():
+        print(name)
+        if name.find('model_ext'):
+            param.requires_grad = False
+        else:
+            param.requires_grad = True
+
+
+def make_class_layer(num_ftrs, num_class, drop_rate=0.2):
+    return nn.Sequential(nn.Linear(num_ftrs, num_class), nn.Dropout(drop_rate), )
+
+
+def get_num_feats(model_name, model_ft, input_tensor):
+    res = ext_features(model_name, model_ft, input_tensor)
+    num_feats = res.shape[-1]
+    return num_feats
+
+def ext_features(model_ft, input_tensor):
+    model_ft.head = nn.Identity()
+    temp = model_ft(input_tensor)
+    res = temp.mean((2,3), keepdim=False)
+    #t_max = temp.amax((2,3), keepdim=False)
+    #print(t_max.shape)
+    #res = torch.cat((t_avg, t_max), -1)
+    return res
+
+
+def make_fe_from_classification_model(model_name, model_ft, num_class, drop_rate=0.2):
+    """
+    :param model_name:
+    :param model_ft:
+    :param num_class:
+    :return:
+    """
+
+    if model_name == 'coca':
+        num_ftrs = 768
+        model_ft = MLP(num_ftrs, num_class, drop=drop_rate, use_pooling=False,
+                       use_ext=True, model_ext=model_ft)
+        set_parameter_requires_grad(model_ft)
+        a = 0
+
+    elif model_name == 'inception_next_tiny' or model_name == 'inception_next_base_384' or model_name == 'inception_next_base':
+        num_ftrs = model_ft.num_features
+        set_parameter_requires_grad_old(model_ft)
+        model_ft.head = MLP(num_ftrs, num_class, drop=drop_rate, use_pooling=True)
+
+    elif model_name == 'edgenext_xx_small_bn_hs' or model_name == 'edgenext_small_usi' or model_name == 'edgenext_base_usi':
+        #set_parameter_requires_grad(model_ft, feat_ext)
+        num_ftrs = model_ft.head.in_features
+        model_ft.head = MLP(num_ftrs, num_class, drop=drop_rate, use_pooling=False)
+    return model_ft
+
+
+def test():
+    from src.utils.utils_img import tensor_from_img
+    num_class = 10
+    path2cfg = fr'{get_proj_dir()}in/config_files/'
+    nn_conf = get_nn_cfg(path2cfg)
+    device = nn_conf.device
+    model = load_models(nn_conf).to(device)
+    print(model)
+    model.eval()
+    img_tensor = torch.rand((16, 3, 224, 224))
+    img_tensor = img_tensor.to(device)
+    #print(img_tensor.shape)
+    #res = ext_features(model, img_tensor)
+    #res = model(img_tensor)
+    #print(res.shape)
+    #print(res)
+    model = make_fe_from_classification_model(nn_conf.model_name, model, num_class)
+    a = 0
+    set_parameter_requires_grad(model)
+    #model.to('cuda')
+
+
+if __name__ == "__main__":
+    test()

src/train/criterion/__init__.py

@@ -0,0 +1 @@
+from src.train.criterion.init_loss import get_loss_func

src/train/criterion/init_loss.py

@@ -0,0 +1,37 @@
+from src.train.criterion.loss import *
+
+
+'''
+assym_ml    --> AsymmetricLossMultiLabel
+assym_sl    --> AsymmetricLossSingleLabel
+bce --> BinaryCrossEntropy
+lab_smooth_cse  --->    LabelSmoothingCrossEntropy
+cse --->    SoftTargetCrossEntropy
+jsd_cse --->    JsdCrossEntropy
+'''
+arr_loss_type = ['assym_ml', 'assym_sl', 'bce', 'cse', 'soft_cse',  'jsd_cse']
+
+
+def get_loss_func(loss_type, label_smooth_coef=0.2, eps=1e-8):
+    """
+    :param loss_type:
+    :param label_smooth_coef:
+    :param eps:
+    :return:
+    """
+    assert loss_type in arr_loss_type, f'loss_type should be {arr_loss_type}'
+
+    loss_func = None
+    if loss_type == 'assym_ml':
+        loss_func = AsymmetricLossMultiLabel(eps=eps)
+    elif loss_type == 'assym_sl':
+        loss_func = AsymmetricLossSingleLabel(eps=eps)
+    elif loss_type == 'bce':
+        loss_func = BinaryCrossEntropy(smoothing=label_smooth_coef)
+    elif loss_type == 'cse':
+        loss_func = LabelSmoothingCrossEntropy(smoothing=label_smooth_coef)
+    elif loss_type == 'soft_cse':
+        loss_func = SoftTargetCrossEntropy()
+    elif loss_type == 'jsd_cse':
+        loss_func = JsdCrossEntropy(smoothing=label_smooth_coef)
+    return loss_func

src/train/criterion/loss/__init__.py

@@ -0,0 +1,4 @@
+from .asymmetric_loss import AsymmetricLossMultiLabel, AsymmetricLossSingleLabel
+from .binary_cross_entropy import BinaryCrossEntropy
+from .cross_entropy import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
+from .jsd import JsdCrossEntropy

src/train/criterion/loss/asymmetric_loss.py

@@ -0,0 +1,97 @@
+import torch
+import torch.nn as nn
+
+
+class AsymmetricLossMultiLabel(nn.Module):
+    def __init__(self, gamma_neg=4, gamma_pos=1, clip=0.05, eps=1e-8, disable_torch_grad_focal_loss=False):
+        super(AsymmetricLossMultiLabel, self).__init__()
+
+        self.gamma_neg = gamma_neg
+        self.gamma_pos = gamma_pos
+        self.clip = clip
+        self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss
+        self.eps = eps
+
+    def forward(self, x, y):
+        """"
+        Parameters
+        ----------
+        x: input logits
+        y: targets (multi-label binarized vector)
+        """
+
+        # Calculating Probabilities
+        x_sigmoid = torch.sigmoid(x)
+        xs_pos = x_sigmoid
+        xs_neg = 1 - x_sigmoid
+
+        # Asymmetric Clipping
+        if self.clip is not None and self.clip > 0:
+            xs_neg = (xs_neg + self.clip).clamp(max=1)
+
+        # Basic CE calculation
+        los_pos = y * torch.log(xs_pos.clamp(min=self.eps))
+        los_neg = (1 - y) * torch.log(xs_neg.clamp(min=self.eps))
+        loss = los_pos + los_neg
+
+        # Asymmetric Focusing
+        if self.gamma_neg > 0 or self.gamma_pos > 0:
+            if self.disable_torch_grad_focal_loss:
+                torch._C.set_grad_enabled(False)
+            pt0 = xs_pos * y
+            pt1 = xs_neg * (1 - y)  # pt = p if t > 0 else 1-p
+            pt = pt0 + pt1
+            one_sided_gamma = self.gamma_pos * y + self.gamma_neg * (1 - y)
+            one_sided_w = torch.pow(1 - pt, one_sided_gamma)
+            if self.disable_torch_grad_focal_loss:
+                torch._C.set_grad_enabled(True)
+            loss *= one_sided_w
+
+        return -loss.sum()
+
+
+class AsymmetricLossSingleLabel(nn.Module):
+    def __init__(self, gamma_pos=1, gamma_neg=4, eps: float = 0.1, reduction='mean'):
+        super(AsymmetricLossSingleLabel, self).__init__()
+
+        self.eps = eps
+        self.logsoftmax = nn.LogSoftmax(dim=-1)
+        self.targets_classes = []  # prevent gpu repeated memory allocation
+        self.gamma_pos = gamma_pos
+        self.gamma_neg = gamma_neg
+        self.reduction = reduction
+
+    def forward(self, inputs, target, reduction=None):
+        """"
+        Parameters
+        ----------
+        x: input logits
+        y: targets (1-hot vector)
+        """
+
+        num_classes = inputs.size()[-1]
+        log_preds = self.logsoftmax(inputs)
+        self.targets_classes = torch.zeros_like(inputs).scatter_(1, target.long().unsqueeze(1), 1)
+
+        # ASL weights
+        targets = self.targets_classes
+        anti_targets = 1 - targets
+        xs_pos = torch.exp(log_preds)
+        xs_neg = 1 - xs_pos
+        xs_pos = xs_pos * targets
+        xs_neg = xs_neg * anti_targets
+        asymmetric_w = torch.pow(1 - xs_pos - xs_neg,
+                                 self.gamma_pos * targets + self.gamma_neg * anti_targets)
+        log_preds = log_preds * asymmetric_w
+
+        if self.eps > 0:  # label smoothing
+            self.targets_classes.mul_(1 - self.eps).add_(self.eps / num_classes)
+
+        # loss calculation
+        loss = - self.targets_classes.mul(log_preds)
+
+        loss = loss.sum(dim=-1)
+        if self.reduction == 'mean':
+            loss = loss.mean()
+
+        return loss

src/train/criterion/loss/binary_cross_entropy.py

@@ -0,0 +1,47 @@
+""" Binary Cross Entropy w/ a few extras
+
+Hacked together by / Copyright 2021 Ross Wightman
+"""
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BinaryCrossEntropy(nn.Module):
+    """ BCE with optional one-hot from dense targets, label smoothing, thresholding
+    NOTE for experiments comparing CE to BCE /w label smoothing, may remove
+    """
+    def __init__(
+            self, smoothing=0.1, target_threshold: Optional[float] = None, weight: Optional[torch.Tensor] = None,
+            reduction: str = 'mean', pos_weight: Optional[torch.Tensor] = None):
+        super(BinaryCrossEntropy, self).__init__()
+        assert 0. <= smoothing < 1.0
+        self.smoothing = smoothing
+        self.target_threshold = target_threshold
+        self.reduction = reduction
+        self.register_buffer('weight', weight)
+        self.register_buffer('pos_weight', pos_weight)
+
+    def forward(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        assert x.shape[0] == target.shape[0]
+        if target.shape != x.shape:
+            # NOTE currently assume smoothing or other label softening is applied upstream if targets are already sparse
+            num_classes = x.shape[-1]
+            # FIXME should off/on be different for smoothing w/ BCE? Other impl out there differ
+            off_value = self.smoothing / num_classes
+            on_value = 1. - self.smoothing + off_value
+            target = target.long().view(-1, 1)
+            target = torch.full(
+                (target.size()[0], num_classes),
+                off_value,
+                device=x.device, dtype=x.dtype).scatter_(1, target, on_value)
+        if self.target_threshold is not None:
+            # Make target 0, or 1 if threshold set
+            target = target.gt(self.target_threshold).to(dtype=target.dtype)
+        return F.binary_cross_entropy_with_logits(
+            x, target,
+            self.weight,
+            pos_weight=self.pos_weight,
+            reduction=self.reduction)

src/train/criterion/loss/cross_entropy.py

@@ -0,0 +1,36 @@
+""" Cross Entropy w/ smoothing or soft targets
+
+Hacked together by / Copyright 2021 Ross Wightman
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LabelSmoothingCrossEntropy(nn.Module):
+    """ NLL loss with label smoothing.
+    """
+    def __init__(self, smoothing=0.1):
+        super(LabelSmoothingCrossEntropy, self).__init__()
+        assert smoothing < 1.0
+        self.smoothing = smoothing
+        self.confidence = 1. - smoothing
+
+    def forward(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        logprobs = F.log_softmax(x, dim=-1)
+        nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1))
+        nll_loss = nll_loss.squeeze(1)
+        smooth_loss = -logprobs.mean(dim=-1)
+        loss = self.confidence * nll_loss + self.smoothing * smooth_loss
+        return loss.mean()
+
+
+class SoftTargetCrossEntropy(nn.Module):
+
+    def __init__(self):
+        super(SoftTargetCrossEntropy, self).__init__()
+
+    def forward(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        loss = torch.sum(-target * F.log_softmax(x, dim=-1), dim=-1)
+        return loss.mean()

src/train/criterion/loss/distillation_loss.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+"""
+Implements the knowledge distillation loss
+"""
+import torch
+from torch.nn import functional as F
+
+class DistillationLoss(torch.nn.Module):
+    """
+    This module wraps a standard criterion and adds an extra knowledge distillation loss by
+    taking a teacher model prediction and using it as additional supervision.
+    """
+
+    def __init__(self, base_criterion: torch.nn.Module, teacher_model: torch.nn.Module,
+                 distillation_type: str, alpha: float, tau: float):
+        super().__init__()
+        self.base_criterion = base_criterion
+        self.teacher_model = teacher_model
+        assert distillation_type in ['none', 'soft', 'hard']
+        self.distillation_type = distillation_type
+        self.alpha = alpha
+        self.tau = tau
+
+    def forward(self, inputs, outputs, labels):
+        """
+        Args:
+            inputs: The original inputs that are feed to the teacher model
+            outputs: the outputs of the model to be trained. It is expected to be
+                either a Tensor, or a Tuple[Tensor, Tensor], with the original output
+                in the first position and the distillation predictions as the second output
+            labels: the labels for the base criterion
+        """
+        outputs_kd = None
+        if not isinstance(outputs, torch.Tensor):
+            # assume that the model outputs a tuple of [outputs, outputs_kd]
+            outputs, outputs_kd = outputs
+        base_loss = self.base_criterion(outputs, labels)
+        if self.distillation_type == 'none':
+            return base_loss
+
+        if outputs_kd is None:
+            raise ValueError("When knowledge distillation is enabled, the model is "
+                             "expected to return a Tuple[Tensor, Tensor] with the output of the "
+                             "class_token and the dist_token")
+        # don't backprop throught the teacher
+        with torch.no_grad():
+            teacher_outputs = self.teacher_model(inputs)
+
+        if self.distillation_type == 'soft':
+            T = self.tau
+            # taken from https://github.com/peterliht/knowledge-distillation-pytorch/blob/master/model/net.py#L100
+            # with slight modifications
+            distillation_loss = F.kl_div(
+                F.log_softmax(outputs_kd / T, dim=1),
+                F.log_softmax(teacher_outputs / T, dim=1),
+                reduction='sum',
+                log_target=True
+            ) * (T * T) / outputs_kd.numel()
+        elif self.distillation_type == 'hard':
+            distillation_loss = F.cross_entropy(
+                outputs_kd, teacher_outputs.argmax(dim=1))
+
+        loss = base_loss * (1 - self.alpha) + distillation_loss * self.alpha
+        return loss

src/train/criterion/loss/jsd.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .cross_entropy import LabelSmoothingCrossEntropy
+
+
+class JsdCrossEntropy(nn.Module):
+    """ Jensen-Shannon Divergence + Cross-Entropy Loss
+
+    Based on impl here: https://github.com/google-research/augmix/blob/master/imagenet.py
+    From paper: 'AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty -
+    https://arxiv.org/abs/1912.02781
+
+    Hacked together by / Copyright 2020 Ross Wightman
+    """
+    def __init__(self, num_splits=3, alpha=12, smoothing=0.1):
+        super().__init__()
+        self.num_splits = num_splits
+        self.alpha = alpha
+        if smoothing is not None and smoothing > 0:
+            self.cross_entropy_loss = LabelSmoothingCrossEntropy(smoothing)
+        else:
+            self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
+
+    def __call__(self, output, target):
+        split_size = output.shape[0] // self.num_splits
+        assert split_size * self.num_splits == output.shape[0]
+        logits_split = torch.split(output, split_size)
+
+        # Cross-entropy is only computed on clean images
+        loss = self.cross_entropy_loss(logits_split[0], target[:split_size])
+        probs = [F.softmax(logits, dim=1) for logits in logits_split]
+
+        # Clamp mixture distribution to avoid exploding KL divergence
+        logp_mixture = torch.clamp(torch.stack(probs).mean(axis=0), 1e-7, 1).log()
+        loss += self.alpha * sum([F.kl_div(
+            logp_mixture, p_split, reduction='batchmean') for p_split in probs]) / len(probs)
+        return loss

src/train/criterion/metrics/accuracy.py

@@ -0,0 +1,19 @@
+import torch
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    with torch.inference_mode():
+        maxk = max(topk)
+        batch_size = target.size(0)
+        if target.ndim == 2:
+            target = target.max(dim=1)[1]
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target[None])
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].flatten().sum(dtype=torch.float32)
+            res.append(correct_k * (100.0 / batch_size))
+        return res

src/train/lr_scheduler_v2/knee_scheduler.py

@@ -0,0 +1,36 @@
+import torch
+from torch.optim.optimizer import Optimizer
+
+class KneeLRScheduler:
+
+    def __init__(self, optimizer, peak_lr, warmup_steps=0, explore_steps=0, total_steps=0):
+        self.optimizer = optimizer
+        self.peak_lr = peak_lr
+        self.warmup_steps = warmup_steps
+        self.explore_steps = (int (total_steps-warmup_steps) // 2 ) if explore_steps == 0 else explore_steps
+        self.total_steps = total_steps
+        self.decay_steps = self.total_steps - (self.explore_steps + self.warmup_steps)
+        self.current_step = 1
+
+        assert self.decay_steps >= 0
+
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = self.get_lr(self.current_step)
+
+        if not isinstance(self.optimizer, Optimizer):
+            raise TypeError('{} is not an Optimizer'.format(
+                type(self.optimizer).__name__))
+
+    def get_lr(self, global_step):
+        if global_step <= self.warmup_steps:
+            return self.peak_lr * global_step / self.warmup_steps
+        elif global_step <= (self.explore_steps + self.warmup_steps):
+            return self.peak_lr
+        else:
+            slope = -1 * self.peak_lr / self.decay_steps
+            return max(0.0, self.peak_lr + slope*(global_step - (self.explore_steps + self.warmup_steps)))
+
+    def step(self):
+        self.current_step += 1
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = self.get_lr(self.current_step)

src/train/lr_scheduler_v2/lr_scheduler.py

@@ -0,0 +1,46 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from torch.optim.lr_scheduler import _LRScheduler
+
+class LearningRateScheduler(_LRScheduler):
+    r"""
+    Provides inteface of learning rate scheduler.
+
+    Note:
+        Do not use this class directly, use one of the sub classes.
+    """
+    def __init__(self, optimizer, lr):
+        self.optimizer = optimizer
+        self.lr = lr
+
+    def step(self, *args, **kwargs):
+        raise NotImplementedError
+
+    @staticmethod
+    def set_lr(optimizer, lr):
+        for g in optimizer.param_groups:
+            g['lr'] = lr
+
+    def get_lr(self):
+        for g in self.optimizer.param_groups:
+            return g['lr']

src/train/lr_scheduler_v2/reduce_lr_on_plateau_lr_scheduler.py

@@ -0,0 +1,69 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from omegaconf import DictConfig
+from torch.optim import Optimizer
+
+from src.train.lr_scheduler_v2.lr_scheduler import LearningRateScheduler
+
+
+class ReduceLROnPlateauScheduler(LearningRateScheduler):
+    r"""
+    Reduce learning rate when a metric has stopped improving. Models often benefit from reducing the learning rate by
+    a factor of 2-10 once learning stagnates. This scheduler reads a metrics quantity and if no improvement is seen
+    for a ‘patience’ number of epochs, the learning rate is reduced.
+
+    Args:
+        optimizer (Optimizer): Optimizer.
+        lr (float): Initial learning rate.
+        patience (int): Number of epochs with no improvement after which learning rate will be reduced.
+        factor (float): Factor by which the learning rate will be reduced. new_lr = lr * factor.
+    """
+    def __init__(
+            self,
+            optimizer: Optimizer,
+            lr: float,
+            patience: int = 1,
+            factor: float = 0.3,
+    ) -> None:
+        super(ReduceLROnPlateauScheduler, self).__init__(optimizer, lr)
+        self.lr = lr
+        self.patience = patience
+        self.factor = factor
+        self.val_loss = 100.0
+        self.count = 0
+
+    def step(self, val_loss: float):
+        if val_loss is not None:
+            if self.val_loss < val_loss:
+                self.count += 1
+                self.val_loss = val_loss
+            else:
+                self.count = 0
+                self.val_loss = val_loss
+
+            if self.patience == self.count:
+                self.count = 0
+                self.lr *= self.factor
+                self.set_lr(self.optimizer, self.lr)
+
+        return self.lr

src/train/lr_scheduler_v2/transformer_lr_scheduler.py

@@ -0,0 +1,93 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import math
+import torch
+from typing import Optional
+from torch.optim import Optimizer
+
+from src.train.lr_scheduler_v2.lr_scheduler import LearningRateScheduler
+
+
+class TransformerLRScheduler(LearningRateScheduler):
+    r"""
+    Transformer Learning Rate Scheduler proposed in "Attention Is All You Need"
+
+    Args:
+        optimizer (Optimizer): Optimizer.
+        init_lr (float): Initial learning rate.
+        peak_lr (float): Maximum learning rate.
+        final_lr (float): Final learning rate.
+        final_lr_scale (float): Final learning rate scale
+        warmup_steps (int): Warmup the learning rate linearly for the first N updates
+        decay_steps (int): Steps in decay stages
+    """
+    def __init__(
+            self,
+            optimizer: Optimizer,
+            init_lr: float,
+            peak_lr: float,
+            final_lr: float,
+            final_lr_scale: float,
+            warmup_steps: int,
+            decay_steps: int,
+    ) -> None:
+        assert isinstance(warmup_steps, int), "warmup_steps should be integer type"
+        assert isinstance(decay_steps, int), "total_steps should be integer type"
+
+        super(TransformerLRScheduler, self).__init__(optimizer, init_lr)
+        self.final_lr = final_lr
+        self.peak_lr = peak_lr
+        self.warmup_steps = warmup_steps
+        self.decay_steps = decay_steps
+
+        self.warmup_rate = self.peak_lr / self.warmup_steps
+        self.decay_factor = -math.log(final_lr_scale) / self.decay_steps
+
+        self.init_lr = init_lr
+        self.update_steps = 0
+
+    def _decide_stage(self):
+        if self.update_steps < self.warmup_steps:
+            return 0, self.update_steps
+
+        if self.warmup_steps <= self.update_steps < self.warmup_steps + self.decay_steps:
+            return 1, self.update_steps - self.warmup_steps
+
+        return 2, None
+
+    def step(self, val_loss: Optional[torch.FloatTensor] = None):
+        self.update_steps += 1
+        stage, steps_in_stage = self._decide_stage()
+
+        if stage == 0:
+            self.lr = self.update_steps * self.warmup_rate
+        elif stage == 1:
+            self.lr = self.peak_lr * math.exp(-self.decay_factor * steps_in_stage)
+        elif stage == 2:
+            self.lr = self.final_lr
+        else:
+            raise ValueError("Undefined stage")
+
+        self.set_lr(self.optimizer, self.lr)
+
+        return self.lr

src/train/lr_scheduler_v2/tri_stage_lr_scheduler.py

@@ -0,0 +1,114 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import math
+import torch
+from typing import Optional
+from torch.optim import Optimizer
+
+from src.train.lr_scheduler_v2.lr_scheduler import LearningRateScheduler
+
+
+class TriStageLRScheduler(LearningRateScheduler):
+    r"""
+    Tri-Stage Learning Rate Scheduler. Implement the learning rate scheduler in "SpecAugment"
+
+    Args:
+        optimizer (Optimizer): Optimizer.
+        init_lr (float): Initial learning rate.
+        peak_lr (float): Maximum learning rate.
+        final_lr (float): Final learning rate.
+        init_lr_scale (float): Initial learning rate scale.
+        final_lr_scale (float): Final learning rate scale.
+        warmup_steps (int): Warmup the learning rate linearly for the first N updates.
+        hold_steps (int): Hold the learning rate for the N updates.
+        decay_steps (int): Decay the learning rate linearly for the first N updates.
+        total_steps (int): Total steps in training.
+    """
+    def __init__(
+            self,
+            optimizer: Optimizer,
+            init_lr: float,
+            peak_lr: float,
+            final_lr: float,
+            init_lr_scale: float,
+            final_lr_scale: float,
+            warmup_steps: int,
+            hold_steps: int,
+            decay_steps: int,
+            total_steps: int,
+    ):
+        assert isinstance(warmup_steps, int), "warmup_steps should be inteager type"
+        assert isinstance(total_steps, int), "total_steps should be inteager type"
+
+        super(TriStageLRScheduler, self).__init__(optimizer, init_lr)
+        self.init_lr = init_lr
+        self.init_lr *= init_lr_scale
+        self.final_lr = final_lr
+        self.peak_lr = peak_lr
+        self.warmup_steps = warmup_steps
+        self.hold_steps = hold_steps
+        self.decay_steps = decay_steps
+
+        self.warmup_rate = (self.peak_lr - self.init_lr) / self.warmup_steps if self.warmup_steps != 0 else 0
+        self.decay_factor = -math.log(final_lr_scale) / self.decay_steps
+
+        self.lr = self.init_lr
+        self.update_steps = 0
+
+    def _decide_stage(self):
+        if self.update_steps < self.warmup_steps:
+            return 0, self.update_steps
+
+        offset = self.warmup_steps
+
+        if self.update_steps < offset + self.hold_steps:
+            return 1, self.update_steps - offset
+
+        offset += self.hold_steps
+
+        if self.update_steps <= offset + self.decay_steps:
+            # decay stage
+            return 2, self.update_steps - offset
+
+        offset += self.decay_steps
+
+        return 3, self.update_steps - offset
+
+    def step(self, val_loss: Optional[torch.FloatTensor] = None):
+        stage, steps_in_stage = self._decide_stage()
+
+        if stage == 0:
+            self.lr = self.init_lr + self.warmup_rate * steps_in_stage
+        elif stage == 1:
+            self.lr = self.peak_lr
+        elif stage == 2:
+            self.lr = self.peak_lr * math.exp(-self.decay_factor * steps_in_stage)
+        elif stage == 3:
+            self.lr = self.final_lr
+        else:
+            raise ValueError("Undefined stage")
+
+        self.set_lr(self.optimizer, self.lr)
+        self.update_steps += 1
+
+        return self.lr

src/train/lr_scheduler_v2/warmup_lr_scheduler.py

@@ -0,0 +1,62 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch
+from typing import Optional
+from torch.optim import Optimizer
+
+from src.train.lr_scheduler_v2.lr_scheduler import LearningRateScheduler
+
+
+class WarmupLRScheduler(LearningRateScheduler):
+    """
+    Warmup learning rate until `total_steps`
+
+    Args:
+        optimizer (Optimizer): wrapped optimizer.
+
+    """
+    def __init__(
+            self,
+            optimizer: Optimizer,
+            init_lr: float,
+            peak_lr: float,
+            warmup_steps: int,
+    ) -> None:
+        super(WarmupLRScheduler, self).__init__(optimizer, init_lr)
+        self.init_lr = init_lr
+        if warmup_steps != 0:
+            warmup_rate = peak_lr - init_lr
+            self.warmup_rate = warmup_rate / warmup_steps
+        else:
+            self.warmup_rate = 0
+        self.update_steps = 1
+        self.lr = init_lr
+        self.warmup_steps = warmup_steps
+
+    def step(self, val_loss: Optional[torch.FloatTensor] = None):
+        if self.update_steps < self.warmup_steps:
+            lr = self.init_lr + self.warmup_rate * self.update_steps
+            self.set_lr(self.optimizer, lr)
+            self.lr = lr
+        self.update_steps += 1
+        return self.lr

src/train/lr_scheduler_v2/warmup_reduce_lr_on_plateau_scheduler.py

@@ -0,0 +1,88 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from torch.optim import Optimizer
+from typing import Optional
+
+from src.train.lr_scheduler_v2.lr_scheduler import LearningRateScheduler
+from src.train.lr_scheduler_v2.reduce_lr_on_plateau_lr_scheduler import ReduceLROnPlateauScheduler
+from src.train.lr_scheduler_v2.warmup_lr_scheduler import WarmupLRScheduler
+
+
+class WarmupReduceLROnPlateauScheduler(LearningRateScheduler):
+    r"""
+    Warmup learning rate until `warmup_steps` and reduce learning rate on plateau after.
+
+    Args:
+        optimizer (Optimizer): wrapped optimizer.
+        init_lr (float): Initial learning rate.
+        peak_lr (float): Maximum learning rate.
+        warmup_steps (int): Warmup the learning rate linearly for the first N updates.
+        patience (int): Number of epochs with no improvement after which learning rate will be reduced.
+        factor (float): Factor by which the learning rate will be reduced. new_lr = lr * factor.
+    """
+    def __init__(
+            self,
+            optimizer: Optimizer,
+            init_lr: float,
+            peak_lr: float,
+            warmup_steps: int,
+            patience: int = 1,
+            factor: float = 0.3,
+    ) -> None:
+        super(WarmupReduceLROnPlateauScheduler, self).__init__(optimizer, init_lr)
+        self.warmup_steps = warmup_steps
+        self.update_steps = 0
+        self.warmup_rate = (peak_lr - init_lr) / self.warmup_steps \
+            if self.warmup_steps != 0 else 0
+        self.schedulers = [
+            WarmupLRScheduler(
+                optimizer=optimizer,
+                init_lr=init_lr,
+                peak_lr=peak_lr,
+                warmup_steps=warmup_steps,
+            ),
+            ReduceLROnPlateauScheduler(
+                optimizer=optimizer,
+                lr=peak_lr,
+                patience=patience,
+                factor=factor,
+            ),
+        ]
+
+    def _decide_stage(self):
+        if self.update_steps < self.warmup_steps:
+            return 0, self.update_steps
+        else:
+            return 1, None
+
+    def step(self, val_loss: Optional[float] = None):
+        stage, steps_in_stage = self._decide_stage()
+
+        if stage == 0:
+            self.schedulers[0].step()
+        elif stage == 1:
+            self.schedulers[1].step(val_loss)
+
+        self.update_steps += 1
+
+        return self.get_lr()

src/train/opt/__init__.py

@@ -0,0 +1 @@
+from src.train.opt.init_opt import create_optimizer

src/train/opt/init_opt.py

@@ -0,0 +1,71 @@
+# https://neptune.ai/blog/how-to-choose-a-learning-rate-scheduler
+# https://d2l.ai/chapter_optimization/lr-scheduler.html
+# https://towardsdatascience.com/the-best-learning-rate-schedules-6b7b9fb72565
+# https://towardsdatascience.com/a-visual-guide-to-learning-rate-schedulers-in-pytorch-24bbb262c863
+# https://huggingface.co/docs/timm/reference/schedulers
+# https://github.com/Tony-Y/pytorch_warmup
+
+# https://jdhao.github.io/2020/08/14/warmup_maskrcnn_how_does_it_work/
+# https://datascience.stackexchange.com/questions/55991/in-the-context-of-deep-learning-what-is-training-warmup-steps
+# https://stackabuse.com/learning-rate-warmup-with-cosine-decay-in-keras-and-tensorflow/
+# https://stackoverflow.com/questions/55933867/what-does-learning-rate-warm-up-mean
+# https://www.programmersought.com/article/76184871189/
+# https://daydreamatnight.github.io/2022/03/08/learning-rate-schedule/
+# https://www.programmersought.com/article/76184871189/
+
+
+from src.train.lr_scheduler_v2.warmup_reduce_lr_on_plateau_scheduler import WarmupReduceLROnPlateauScheduler
+from src.train.opt.opt_modules import *
+import torch_optimizer as optim
+
+arr_opt = ['lamb', 'lion', 'adamw', 'adam', 'adafactor', 'ranger']
+arr_lr_sch = ['reducelr_plateau']
+
+def create_optimizer(model, warmup_steps,
+                     initial_lr=0.000176, max_lr=1e-4, lr_decay_rate=0.3, weight_decay_rate=0.01,
+                     opt_beta1=0.9, opt_beta2= 0.999, opt_eps=1e-6,
+                     opt_type='lamb', lr_schedule_type='reducelr_plateau',
+                     momentum=0.9, use_lookahead_opt=True):
+    """
+    Args:
+       initial_lr:
+       num_train_steps:
+       num_warmup_steps:
+       weight_decay_rate:
+       opt_type:
+       opt_beta1:
+       opt_beta2:
+       opt_epsilon:
+       momentum
+       use_lookahead_opt
+    Returns:
+    """
+    assert opt_type in arr_opt, f"opt_type should be {arr_opt}"
+    assert lr_schedule_type in lr_schedule_type, f"lr_schedule_type should be {arr_lr_sch}"
+    opt_type = 'adam' if opt_type in ['lamb', 'adamw', 'lion', 'radam'] and weight_decay_rate == 0 else opt_type
+    optimizer, lr_scheduler = None, None
+
+    '''if opt_type == 'lion':
+        optimizer = Lion(model.parameters(),
+                         lr=initial_lr, betas=(opt_beta1, opt_beta2), weight_decay=weight_decay_rate,
+                         use_triton=False)
+    elif opt_type == 'lamb':
+        optimizer = Lamb(model.parameters(),lr=initial_lr, betas=(opt_beta1, opt_beta2), eps=opt_eps,
+                         weight_decay=weight_decay_rate)
+    elif opt_type == 'adamw':
+        optimizer = AdamW(model.parameters(),lr=initial_lr, betas=(opt_beta1, opt_beta2), eps=opt_eps,
+                          weight_decay=weight_decay_rate)
+
+    if use_lookahead_opt: optimizer = Lookahead(optimizer)
+    '''
+
+    optimizer = optim.Ranger(model.parameters(), lr=initial_lr, alpha=0.5, k=6, N_sma_threshhold=5,
+                             betas=(opt_beta1, opt_beta2), eps=opt_eps, weight_decay=weight_decay_rate
+    )
+
+
+    if lr_schedule_type == 'reducelr_plateau':
+        lr_scheduler = WarmupReduceLROnPlateauScheduler(optimizer, init_lr=initial_lr, peak_lr=max_lr,
+                                                        warmup_steps=warmup_steps, patience=1, factor=lr_decay_rate,
+                                                        )
+    return optimizer, lr_scheduler

src/train/opt/opt_modules/__init__.py

@@ -0,0 +1,6 @@
+from torch.optim import Adam, AdamW, NAdam, RMSprop, ASGD, Adamax, Adadelta, Adagrad, RAdam
+from .adafactor import Adafactor
+from .lamb_opt import Lamb
+from .lion_opt import Lion
+from .lookahead_opt import Lookahead
+from .radam_opt import RAdam

src/train/opt/opt_modules/adafactor.py

@@ -0,0 +1,240 @@
+from torch.optim import Optimizer
+
+class Adafactor(Optimizer):
+    """
+    AdaFactor pytorch implementation can be used as a drop in replacement for Adam original fairseq code:
+    https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py
+
+    Paper: *Adafactor: Adaptive Learning Rates with Sublinear Memory Cost* https://arxiv.org/abs/1804.04235 Note that
+    this optimizer internally adjusts the learning rate depending on the `scale_parameter`, `relative_step` and
+    `warmup_init` options. To use a manual (external) learning rate schedule you should set `scale_parameter=False` and
+    `relative_step=False`.
+
+    Arguments:
+        params (`Iterable[nn.parameter.Parameter]`):
+            Iterable of parameters to optimize or dictionaries defining parameter groups.
+        lr (`float`, *optional*):
+            The external learning rate.
+        eps (`Tuple[float, float]`, *optional*, defaults to (1e-30, 1e-3)):
+            Regularization constants for square gradient and parameter scale respectively
+        clip_threshold (`float`, *optional*, defaults 1.0):
+            Threshold of root mean square of final gradient update
+        decay_rate (`float`, *optional*, defaults to -0.8):
+            Coefficient used to compute running averages of square
+        beta1 (`float`, *optional*):
+            Coefficient used for computing running averages of gradient
+        weight_decay (`float`, *optional*, defaults to 0):
+            Weight decay (L2 penalty)
+        scale_parameter (`bool`, *optional*, defaults to `True`):
+            If True, learning rate is scaled by root mean square
+        relative_step (`bool`, *optional*, defaults to `True`):
+            If True, time-dependent learning rate is computed instead of external learning rate
+        warmup_init (`bool`, *optional*, defaults to `False`):
+            Time-dependent learning rate computation depends on whether warm-up initialization is being used
+
+    This implementation handles low-precision (FP16, bfloat) values, but we have not thoroughly tested.
+
+    Recommended T5 finetuning settings (https://discuss.huggingface.co/t/t5-finetuning-tips/684/3):
+
+        - Training without LR warmup or clip_threshold is not recommended.
+
+           - use scheduled LR warm-up to fixed LR
+           - use clip_threshold=1.0 (https://arxiv.org/abs/1804.04235)
+        - Disable relative updates
+        - Use scale_parameter=False
+        - Additional optimizer operations like gradient clipping should not be used alongside Adafactor
+
+    Example:
+
+    ```python
+    Adafactor(model.parameters(), scale_parameter=False, relative_step=False, warmup_init=False, lr=1e-3)
+    ```
+
+    Others reported the following combination to work well:
+
+    ```python
+    Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
+    ```
+
+    When using `lr=None` with [`Trainer`] you will most likely need to use [`~optimization.AdafactorSchedule`]
+    scheduler as following:
+
+    ```python
+    from transformers.optimization import Adafactor, AdafactorSchedule
+
+    optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
+    lr_scheduler = AdafactorSchedule(optimizer)
+    trainer = Trainer(..., optimizers=(optimizer, lr_scheduler))
+    ```
+
+    Usage:
+
+    ```python
+    # replace AdamW with Adafactor
+    optimizer = Adafactor(
+        model.parameters(),
+        lr=1e-3,
+        eps=(1e-30, 1e-3),
+        clip_threshold=1.0,
+        decay_rate=-0.8,
+        beta1=None,
+        weight_decay=0.0,
+        relative_step=False,
+        scale_parameter=False,
+        warmup_init=False,
+    )
+    ```"""
+
+    def __init__(
+            self,
+            params,
+            lr=None,
+            eps=(1e-30, 1e-3),
+            clip_threshold=1.0,
+            decay_rate=-0.8,
+            beta1=None,
+            weight_decay=0.0,
+            scale_parameter=True,
+            relative_step=True,
+            warmup_init=False,
+    ):
+        require_version("torch>=1.5.0")  # add_ with alpha
+        if lr is not None and relative_step:
+            raise ValueError("Cannot combine manual `lr` and `relative_step=True` options")
+        if warmup_init and not relative_step:
+            raise ValueError("`warmup_init=True` requires `relative_step=True`")
+
+        defaults = dict(
+            lr=lr,
+            eps=eps,
+            clip_threshold=clip_threshold,
+            decay_rate=decay_rate,
+            beta1=beta1,
+            weight_decay=weight_decay,
+            scale_parameter=scale_parameter,
+            relative_step=relative_step,
+            warmup_init=warmup_init,
+        )
+        super().__init__(params, defaults)
+
+    @staticmethod
+    def _get_lr(param_group, param_state):
+        rel_step_sz = param_group["lr"]
+        if param_group["relative_step"]:
+            min_step = 1e-6 * param_state["step"] if param_group["warmup_init"] else 1e-2
+            rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"]))
+        param_scale = 1.0
+        if param_group["scale_parameter"]:
+            param_scale = max(param_group["eps"][1], param_state["RMS"])
+        return param_scale * rel_step_sz
+
+    @staticmethod
+    def _get_options(param_group, param_shape):
+        factored = len(param_shape) >= 2
+        use_first_moment = param_group["beta1"] is not None
+        return factored, use_first_moment
+
+    @staticmethod
+    def _rms(tensor):
+        return tensor.norm(2) / (tensor.numel() ** 0.5)
+
+    @staticmethod
+    def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):
+        # copy from fairseq's adafactor implementation:
+        # https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505
+        r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
+        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
+        return torch.mul(r_factor, c_factor)
+
+    def step(self, closure=None):
+        """
+        Performs a single optimization step
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
+                if grad.is_sparse:
+                    raise RuntimeError("Adafactor does not support sparse gradients.")
+
+                state = self.state[p]
+                grad_shape = grad.shape
+
+                factored, use_first_moment = self._get_options(group, grad_shape)
+                # State Initialization
+                if len(state) == 0:
+                    state["step"] = 0
+
+                    if use_first_moment:
+                        # Exponential moving average of gradient values
+                        state["exp_avg"] = torch.zeros_like(grad)
+                    if factored:
+                        state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad)
+                        state["exp_avg_sq_col"] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)
+                    else:
+                        state["exp_avg_sq"] = torch.zeros_like(grad)
+
+                    state["RMS"] = 0
+                else:
+                    if use_first_moment:
+                        state["exp_avg"] = state["exp_avg"].to(grad)
+                    if factored:
+                        state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)
+                        state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)
+                    else:
+                        state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)
+
+                p_data_fp32 = p.data
+                if p.data.dtype in {torch.float16, torch.bfloat16}:
+                    p_data_fp32 = p_data_fp32.float()
+
+                state["step"] += 1
+                state["RMS"] = self._rms(p_data_fp32)
+                lr = self._get_lr(group, state)
+
+                beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])
+                update = (grad**2) + group["eps"][0]
+                if factored:
+                    exp_avg_sq_row = state["exp_avg_sq_row"]
+                    exp_avg_sq_col = state["exp_avg_sq_col"]
+
+                    exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t))
+                    exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t))
+
+                    # Approximation of exponential moving average of square of gradient
+                    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
+                    update.mul_(grad)
+                else:
+                    exp_avg_sq = state["exp_avg_sq"]
+
+                    exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t))
+                    update = exp_avg_sq.rsqrt().mul_(grad)
+
+                update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
+                update.mul_(lr)
+
+                if use_first_moment:
+                    exp_avg = state["exp_avg"]
+                    exp_avg.mul_(group["beta1"]).add_(update, alpha=(1 - group["beta1"]))
+                    update = exp_avg
+
+                if group["weight_decay"] != 0:
+                    p_data_fp32.add_(p_data_fp32, alpha=(-group["weight_decay"] * lr))
+
+                p_data_fp32.add_(-update)
+
+                if p.data.dtype in {torch.float16, torch.bfloat16}:
+                    p.data.copy_(p_data_fp32)
+
+        return loss

src/train/opt/opt_modules/base_optimiser.py

@@ -0,0 +1,388 @@
+import math
+from abc import ABC, abstractmethod
+
+from typing import Callable, Dict, Iterable, Literal, Optional, Tuple, Type, Union, List
+
+import torch
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+
+CLOSURE = Optional[Callable[[], float]]
+LOSS = Optional[float]
+BETAS = Union[Tuple[float, float], Tuple[float, float, float], Tuple[None, float]]
+DEFAULTS = Dict
+PARAMETERS = Optional[Union[Iterable[Dict], Iterable[torch.Tensor]]]
+STATE = Dict
+OPTIMIZER = Type[Optimizer]
+SCHEDULER = Type[LRScheduler]
+
+HUTCHINSON_G = Literal['gaussian', 'rademacher']
+CLASS_MODE = Literal['binary', 'multiclass', 'multilabel']
+
+
+class NoSparseGradientError(Exception):
+    """Raised when the gradient is sparse gradient.
+
+    :param optimizer_name: str. optimizer name.
+    :param note: str. special conditions to note (default '').
+    """
+
+    def __init__(self, optimizer_name: str, note: str = ''):
+        self.note: str = ' ' if not note else f' w/ {note} '
+        self.message: str = f'[-] {optimizer_name}{self.note}does not support sparse gradient.'
+        super().__init__(self.message)
+
+
+class ZeroParameterSizeError(Exception):
+    """Raised when the parameter size is 0."""
+
+    def __init__(self):
+        self.message: str = '[-] parameter size is 0'
+        super().__init__(self.message)
+
+
+class NoClosureError(Exception):
+    """Raised when there's no closure function."""
+
+    def __init__(self, optimizer_name: str, note: str = ''):
+        self.message: str = f'[-] {optimizer_name} requires closure.{note}'
+        super().__init__(self.message)
+
+
+class NegativeLRError(Exception):
+    """Raised when learning rate is negative."""
+
+    def __init__(self, lr: float, lr_type: str = ''):
+        self.note: str = lr_type if lr_type else 'learning rate'
+        self.message: str = f'[-] {self.note} must be positive. ({lr} > 0)'
+        super().__init__(self.message)
+
+
+class NegativeStepError(Exception):
+    """Raised when step is negative."""
+
+    def __init__(self, num_steps: int, step_type: str = ''):
+        self.note: str = step_type if step_type else 'step'
+        self.message: str = f'[-] {self.note} must be positive. ({num_steps} > 0)'
+        super().__init__(self.message)
+
+class BaseOptimizer(ABC, Optimizer):
+    r"""Base optimizer class. Provides common functionalities for the optimizers."""
+
+    def __init__(self, params: PARAMETERS, defaults: DEFAULTS) -> None:
+        super().__init__(params, defaults)
+
+    @staticmethod
+    @torch.no_grad()
+    def set_hessian(param_groups: PARAMETERS, state: STATE, hessian: List[torch.Tensor]) -> None:
+        r"""Set hessian to state from external source. Generally useful when using functorch as a base.
+
+        Example:
+        -------
+            Here's an example::
+
+                # Hutchinson's Estimator using HVP
+                noise = tree_map(lambda v: torch.randn_like(v), params)
+                loss_, hvp_est = jvp(grad(run_model_fn), (params,), (noise,))
+                hessian_diag_est  = tree_map(lambda a, b: a * b, hvp_est, noise)
+
+                optimizer.set_hessian(hessian_diag_est)
+                # OR
+                optimizer.step(hessian=hessian_diag_est)
+
+        :param param_groups: PARAMETERS. parameter groups.
+        :param state: STATE. optimizer state.
+        :param hessian: List[torch.Tensor]. sequence of hessian to set.
+        """
+        i: int = 0
+        for group in param_groups:
+            for p in group['params']:
+                if p.size() != hessian[i].size():
+                    raise ValueError(
+                        f'[-] the shape of parameter and hessian does not match. {p.size()} vs {hessian[i].size()}'
+                    )
+
+                state[p]['hessian'] = hessian[i]
+                i += 1
+
+    @staticmethod
+    def zero_hessian(param_groups: PARAMETERS, state: STATE, pre_zero: bool = True) -> None:
+        r"""Zero-out hessian.
+
+        :param param_groups: PARAMETERS. parameter groups.
+        :param state: STATE. optimizer state.
+        :param pre_zero: bool. zero-out hessian before computing the hessian.
+        """
+        for group in param_groups:
+            for p in group['params']:
+                if p.requires_grad and p.grad is not None and not p.grad.is_sparse:
+                    if 'hessian' not in state[p]:
+                        state[p]['hessian'] = torch.zeros_like(p)
+                    elif pre_zero:
+                        state[p]['hessian'].zero_()
+
+    @staticmethod
+    @torch.no_grad()
+    def compute_hutchinson_hessian(
+            param_groups: PARAMETERS,
+            state: STATE,
+            num_samples: int = 1,
+            alpha: float = 1.0,
+            distribution: HUTCHINSON_G = 'gaussian',
+    ) -> None:
+        r"""Hutchinson's approximate hessian, added to the state under key `hessian`.
+
+        :param param_groups: PARAMETERS. parameter groups.
+        :param state: STATE. optimizer state.
+        :param num_samples: int. number of times to sample `z` for the approximation of the hessian trace.
+        :param alpha: float. alpha.
+        :param distribution: HUTCHINSON_G. type of distribution.
+        """
+        if distribution not in ('gaussian', 'rademacher'):
+            raise NotImplementedError(f'[-] Hessian with distribution {distribution} is not implemented.')
+
+        params: List[torch.Tensor] = [
+            p
+            for group in param_groups
+            for p in group['params']
+            if p.requires_grad and p.grad is not None and not p.grad.is_sparse
+        ]
+        if len(params) == 0:
+            return
+
+        grads = [p.grad for p in params]
+
+        for i in range(num_samples):
+            if distribution == 'rademacher':
+                zs = [torch.randint_like(p, 0, 1) * 2.0 - 1.0 for p in params]
+            else:
+                zs = [torch.randn_like(p) for p in params]
+
+            h_zs = torch.autograd.grad(grads, params, grad_outputs=zs, retain_graph=i < num_samples - 1)
+            for h_z, z, p in zip(h_zs, zs, params):
+                state[p]['hessian'].add_(h_z * z, alpha=alpha / num_samples)
+
+    @staticmethod
+    def apply_weight_decay(
+            p: torch.Tensor,
+            grad: Optional[torch.Tensor],
+            lr: float,
+            weight_decay: float,
+            weight_decouple: bool,
+            fixed_decay: bool,
+            ratio: Optional[float] = None,
+    ) -> None:
+        r"""Apply weight decay.
+
+        :param p: torch.Tensor. parameter.
+        :param grad: torch.Tensor. gradient.
+        :param lr: float. learning rate.
+        :param weight_decay: float. weight decay (L2 penalty).
+        :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+        :param fixed_decay: bool. fix weight decay.
+        :param ratio: Optional[float]. scale weight decay.
+        """
+        if weight_decouple:
+            p.mul_(1.0 - weight_decay * (1.0 if fixed_decay else lr) * (ratio if ratio is not None else 1.0))
+        elif weight_decay > 0.0 and grad is not None:
+            grad.add_(p, alpha=weight_decay)
+
+    @staticmethod
+    def apply_ams_bound(
+            ams_bound: bool, exp_avg_sq: torch.Tensor, max_exp_avg_sq: Optional[torch.Tensor], eps: float
+    ) -> torch.Tensor:
+        r"""Apply AMSBound variant.
+
+        :param ams_bound: bool. whether to apply AMSBound.
+        :param exp_avg_sq: torch.Tensor. exp_avg_sq.
+        :param max_exp_avg_sq: Optional[torch.Tensor]. max_exp_avg_sq.
+        :param eps: float. epsilon.
+        """
+        if ams_bound:
+            torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+            de_nom = max_exp_avg_sq.add(eps)
+        else:
+            de_nom = exp_avg_sq.add(eps)
+
+        return de_nom.sqrt_().add_(eps)
+
+    @staticmethod
+    def debias(beta: float, step: int) -> float:
+        r"""Adam-style debias correction. Returns `1.0 - beta ** step`.
+
+        :param beta: float. beta.
+        :param step: int. number of step.
+        """
+        return 1.0 - math.pow(beta, step)  # fmt: skip
+
+    @staticmethod
+    def debias_beta(beta: float, step: int) -> float:
+        r"""Apply the Adam-style debias correction into beta.
+
+        Simplified version of `\^{beta} = beta * (1.0 - beta ** (step - 1)) / (1.0 - beta ** step)`
+
+        :param beta: float. beta.
+        :param step: int. number of step.
+        """
+        beta_n: float = math.pow(beta, step)
+        return (beta_n - beta) / (beta_n - 1.0)  # fmt: skip
+
+    @staticmethod
+    def apply_adam_debias(adam_debias: bool, step_size: float, bias_correction1: float) -> float:
+        r"""Apply AdamD variant.
+
+        :param adam_debias: bool. whether to apply AdamD.
+        :param step_size: float. step size.
+        :param bias_correction1: float. bias_correction.
+        """
+        return step_size if adam_debias else step_size / bias_correction1
+
+    @staticmethod
+    def get_rectify_step_size(
+            is_rectify: bool,
+            step: int,
+            lr: float,
+            beta2: float,
+            n_sma_threshold: int,
+            degenerated_to_sgd: bool,
+    ) -> Tuple[float, float]:
+        r"""Get step size for rectify optimizer.
+
+        :param is_rectify: bool. whether to apply rectify-variant.
+        :param step: int. number of steps.
+        :param lr: float. learning rate.
+        :param beta2: float. beta2.
+        :param n_sma_threshold: float. SMA threshold.
+        :param degenerated_to_sgd: bool. degenerated to SGD.
+        """
+        step_size: float = lr
+        n_sma: float = 0.0
+
+        if is_rectify:
+            n_sma_max: float = 2.0 / (1.0 - beta2) - 1.0
+            beta2_t: float = beta2 ** step  # fmt: skip
+            n_sma: float = n_sma_max - 2 * step * beta2_t / (1.0 - beta2_t)
+
+            if n_sma >= n_sma_threshold:
+                rt = math.sqrt(
+                    (1.0 - beta2_t) * (n_sma - 4) / (n_sma_max - 4) * (n_sma - 2) / n_sma * n_sma_max / (n_sma_max - 2)
+                )
+            elif degenerated_to_sgd:
+                rt = 1.0
+            else:
+                rt = -1.0
+
+            step_size *= rt
+
+        return step_size, n_sma
+
+    @staticmethod
+    def get_adanorm_gradient(
+            grad: torch.Tensor, adanorm: bool, exp_grad_norm: Optional[torch.Tensor] = None, r: Optional[float] = 0.95
+    ) -> torch.Tensor:
+        r"""Get AdaNorm gradient.
+
+        :param grad: torch.Tensor. gradient.
+        :param adanorm: bool. whether to apply AdaNorm.
+        :param exp_grad_norm: Optional[torch.Tensor]. exp_grad_norm.
+        :param r: float. Optional[float]. momentum (ratio).
+        """
+        if not adanorm or exp_grad_norm is None:
+            return grad
+
+        grad_norm = torch.linalg.norm(grad)
+
+        exp_grad_norm.mul_(r).add_(grad_norm, alpha=1.0 - r)
+
+        return grad.mul(exp_grad_norm).div_(grad_norm) if exp_grad_norm > grad_norm else grad
+
+    @staticmethod
+    def get_rms(x: torch.Tensor) -> float:
+        r"""Get RMS."""
+        return x.norm(2) / math.sqrt(x.numel())
+
+    @staticmethod
+    def approximate_sq_grad(
+            exp_avg_sq_row: torch.Tensor,
+            exp_avg_sq_col: torch.Tensor,
+            output: torch.Tensor,
+    ) -> None:
+        r"""Get approximation of EMA of squared gradient."""
+        r_factor: torch.Tensor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
+        c_factor: torch.Tensor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
+        torch.mul(r_factor, c_factor, out=output)
+
+    @staticmethod
+    def validate_range(x: float, name: str, low: float, high: float, range_type: str = '[)') -> None:
+        if range_type == '[)' and not low <= x < high:
+            raise ValueError(f'[-] {name} must be in the range [{low}, {high})')
+        if range_type == '[]' and not low <= x <= high:
+            raise ValueError(f'[-] {name} must be in the range [{low}, {high}]')
+        if range_type == '(]' and not low < x <= high:
+            raise ValueError(f'[-] {name} must be in the range ({low}, {high}]')
+        if range_type == '()' and not low < x < high:
+            raise ValueError(f'[-] {name} must be in the range ({low}, {high})')
+
+    @staticmethod
+    def validate_non_negative(x: Optional[float], name: str) -> None:
+        if x is not None and x < 0.0:
+            raise ValueError(f'[-] {name} must be non-negative')
+
+    @staticmethod
+    def validate_positive(x: Union[float, int], name: str) -> None:
+        if x <= 0:
+            raise ValueError(f'[-] {name} must be positive')
+
+    @staticmethod
+    def validate_boundary(constant: float, boundary: float, bound_type: str = 'upper') -> None:
+        if bound_type == 'upper' and constant > boundary:
+            raise ValueError(f'[-] constant {constant} must be in a range of (-inf, {boundary}]')
+        if bound_type == 'lower' and constant < boundary:
+            raise ValueError(f'[-] constant {constant} must be in a range of [{boundary}, inf)')
+
+    @staticmethod
+    def validate_step(step: int, step_type: str) -> None:
+        if step < 1:
+            raise NegativeStepError(step, step_type=step_type)
+
+    @staticmethod
+    def validate_options(x: str, name: str, options: List[str]) -> None:
+        if x not in options:
+            opts: str = ' or '.join([f'\'{option}\'' for option in options]).strip()
+            raise ValueError(f'[-] {name} {x} must be one of ({opts})')
+
+    @staticmethod
+    def validate_learning_rate(learning_rate: Optional[float]) -> None:
+        if learning_rate is not None and learning_rate < 0.0:
+            raise NegativeLRError(learning_rate)
+
+    @staticmethod
+    def validate_mod(x: int, y: int) -> None:
+        if x % y != 0:
+            raise ValueError(f'[-] {x} must be divisible by {y}')
+
+    def validate_betas(self, betas: BETAS) -> None:
+        if betas[0] is not None:
+            self.validate_range(betas[0], 'beta1', 0.0, 1.0, range_type='[]')
+
+        self.validate_range(betas[1], 'beta2', 0.0, 1.0, range_type='[]')
+
+        if len(betas) < 3:
+            return
+
+        if betas[2] is not None:
+            self.validate_range(betas[2], 'beta3', 0.0, 1.0, range_type='[]')
+
+    def validate_nus(self, nus: Union[float, Tuple[float, float]]) -> None:
+        if isinstance(nus, float):
+            self.validate_range(nus, 'nu', 0.0, 1.0, range_type='[]')
+        else:
+            self.validate_range(nus[0], 'nu1', 0.0, 1.0, range_type='[]')
+            self.validate_range(nus[1], 'nu2', 0.0, 1.0, range_type='[]')
+
+    @abstractmethod
+    def reset(self) -> None:  # pragma: no cover
+        raise NotImplementedError
+
+    def step(self, closure: CLOSURE = None) -> LOSS:  # pragma: no cover
+        raise NotImplementedError

src/train/opt/opt_modules/ema.py

@@ -0,0 +1,14 @@
+import torch
+
+class ExponentialMovingAverage(torch.optim.swa_utils.AveragedModel):
+    """Maintains moving averages of model parameters using an exponential decay.
+    ``ema_avg = decay * avg_model_param + (1 - decay) * model_param``
+    `torch.optim.swa_utils.AveragedModel <https://pytorch.org/docs/stable/optim.html#custom-averaging-strategies>`_
+    is used to compute the EMA.
+    """
+
+    def __init__(self, model, decay, device="cpu"):
+        def ema_avg(avg_model_param, model_param, num_averaged):
+            return decay * avg_model_param + (1 - decay) * model_param
+
+        super().__init__(model, device, ema_avg, use_buffers=True)

src/train/opt/opt_modules/lamb_opt.py

@@ -0,0 +1,153 @@
+import math
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+from src.utils.utils_types import Betas2, OptFloat, OptLossClosure, Params
+
+__all__ = ('Lamb',)
+
+
+class Lamb(Optimizer):
+    r"""Implements Lamb algorithm.
+
+    It has been proposed in `Large Batch Optimization for Deep Learning:
+    Training BERT in 76 minutes`__.
+
+    Arguments:
+        params: iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr: learning rate (default: 1e-3)
+        betas: coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps: term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay: weight decay (L2 penalty) (default: 0)
+        clamp_value: clamp weight_norm in (0,clamp_value) (default: 10)
+            set to a high value to avoid it (e.g 10e3)
+        adam: always use trust ratio = 1, which turns this
+            into Adam. Useful for comparison purposes. (default: False)
+        debias: debias adam by (1 - beta**step) (default: False)
+
+    Example:
+        >>> optimizer = optim.Lamb(model.parameters(), lr=0.1)
+        >>> optimizer.zero_grad()
+        >>> loss_fn(model(input), target).backward()
+        >>> optimizer.step()
+
+    __ https://arxiv.org/abs/1904.00962
+
+    Note:
+        Reference code: https://github.com/cybertronai/pytorch-lamb
+    """
+
+    def __init__(
+            self,
+            params: Params,
+            lr: float = 1e-3,
+            betas: Betas2 = (0.9, 0.999),
+            eps: float = 1e-6,
+            weight_decay: float = 0,
+            clamp_value: float = 10,
+            adam: bool = False,
+            debias: bool = False,
+    ) -> None:
+        if lr <= 0.0:
+            raise ValueError('Invalid learning rate: {}'.format(lr))
+        if eps < 0.0:
+            raise ValueError('Invalid epsilon value: {}'.format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(
+                'Invalid beta parameter at index 0: {}'.format(betas[0])
+            )
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(
+                'Invalid beta parameter at index 1: {}'.format(betas[1])
+            )
+        if weight_decay < 0:
+            raise ValueError(
+                'Invalid weight_decay value: {}'.format(weight_decay)
+            )
+        if clamp_value < 0.0:
+            raise ValueError('Invalid clamp value: {}'.format(clamp_value))
+
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        self.clamp_value = clamp_value
+        self.adam = adam
+        self.debias = debias
+
+        super(Lamb, self).__init__(params, defaults)
+
+    def step(self, closure: OptLossClosure = None) -> OptFloat:
+        r"""Performs a single optimization step.
+
+        Arguments:
+            closure: A closure that reevaluates the model and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    msg = (
+                        'Lamb does not support sparse gradients, '
+                        'please consider SparseAdam instead'
+                    )
+                    raise RuntimeError(msg)
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                # Decay the first and second moment running average coefficient
+                # m_t
+                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+                # v_t
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+                # Paper v3 does not use debiasing.
+                if self.debias:
+                    bias_correction = math.sqrt(1 - beta2 ** state['step'])
+                    bias_correction /= 1 - beta1 ** state['step']
+                else:
+                    bias_correction = 1
+
+                # Apply bias to lr to avoid broadcast.
+                step_size = group['lr'] * bias_correction
+
+                weight_norm = torch.norm(p.data).clamp(0, self.clamp_value)
+
+                adam_step = exp_avg / exp_avg_sq.sqrt().add(group['eps'])
+                if group['weight_decay'] != 0:
+                    adam_step.add_(p.data, alpha=group['weight_decay'])
+
+                adam_norm = torch.norm(adam_step)
+                if weight_norm == 0 or adam_norm == 0:
+                    trust_ratio = 1
+                else:
+                    trust_ratio = weight_norm / adam_norm
+                state['weight_norm'] = weight_norm
+                state['adam_norm'] = adam_norm
+                state['trust_ratio'] = trust_ratio
+                if self.adam:
+                    trust_ratio = 1
+
+                p.data.add_(adam_step, alpha=-step_size * trust_ratio)
+
+        return loss

src/train/opt/opt_modules/lion_opt.py

@@ -0,0 +1,88 @@
+from typing import Tuple, Optional, Callable
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+# functions
+
+def exists(val):
+    return val is not None
+
+# update functions
+
+def update_fn(p, grad, exp_avg, lr, wd, beta1, beta2):
+    # stepweight decay
+
+    p.data.mul_(1 - lr * wd)
+
+    # weight update
+
+    update = exp_avg.clone().mul_(beta1).add(grad, alpha = 1 - beta1).sign_()
+    p.add_(update, alpha = -lr)
+
+    # decay the momentum running average coefficient
+
+    exp_avg.mul_(beta2).add_(grad, alpha = 1 - beta2)
+
+# class
+
+class Lion(Optimizer):
+    def __init__(
+            self,
+            params,
+            lr: float = 1e-4,
+            betas: Tuple[float, float] = (0.9, 0.99),
+            weight_decay: float = 0.0,
+            use_triton: bool = False
+    ):
+        assert lr > 0.
+        assert all([0. <= beta <= 1. for beta in betas])
+
+        defaults = dict(
+            lr = lr,
+            betas = betas,
+            weight_decay = weight_decay
+        )
+
+        super().__init__(params, defaults)
+
+        self.update_fn = update_fn
+
+        if use_triton:
+            from src.train.opt.opt_modules.lion_triton_opt import update_fn as triton_update_fn
+            self.update_fn = triton_update_fn
+
+    @torch.no_grad()
+    def step(
+            self,
+            closure: Optional[Callable] = None
+    ):
+
+        loss = None
+        if exists(closure):
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in filter(lambda p: exists(p.grad), group['params']):
+
+                grad, lr, wd, beta1, beta2, state = p.grad, group['lr'], group['weight_decay'], *group['betas'], self.state[p]
+
+                # init state - exponential moving average of gradient values
+
+                if len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p)
+
+                exp_avg = state['exp_avg']
+
+                self.update_fn(
+                    p,
+                    grad,
+                    exp_avg,
+                    lr,
+                    wd,
+                    beta1,
+                    beta2
+                )
+
+        return loss

src/train/opt/opt_modules/lion_triton_opt.py

@@ -0,0 +1,97 @@
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+except ImportError as e:
+    print('triton is not installed, please install by running `pip install triton -U --pre`')
+    exit()
+
+
+@triton.autotune(configs = [
+    triton.Config({'BLOCK_SIZE': 128}, num_warps = 4),
+    triton.Config({'BLOCK_SIZE': 1024}, num_warps = 8),
+], key = ['n_elements'])
+@triton.jit
+def update_fn_kernel(
+        p_ptr,
+        grad_ptr,
+        exp_avg_ptr,
+        lr,
+        wd,
+        beta1,
+        beta2,
+        n_elements,
+        BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis = 0)
+
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+
+    mask = offsets < n_elements
+
+    # offsetted pointers
+
+    offset_p_ptr = p_ptr + offsets
+    offset_grad_ptr = grad_ptr + offsets
+    offset_exp_avg_ptr = exp_avg_ptr + offsets
+
+    # load
+
+    p = tl.load(offset_p_ptr, mask = mask)
+    grad = tl.load(offset_grad_ptr, mask = mask)
+    exp_avg = tl.load(offset_exp_avg_ptr, mask = mask)
+
+    # stepweight decay
+
+    p = p * (1 - lr * wd)
+
+    # diff between momentum running average and grad
+
+    diff = exp_avg - grad
+
+    # weight update
+
+    update = diff * beta1 + grad
+
+    # torch.sign
+
+    can_update = update != 0
+    update_sign = tl.where(update > 0, -lr, lr)
+
+    p = p + update_sign * can_update
+
+    # decay the momentum running average coefficient
+
+    exp_avg = diff * beta2 + grad
+
+    # store new params and momentum running average coefficient
+
+    tl.store(offset_p_ptr, p, mask = mask)
+    tl.store(offset_exp_avg_ptr, exp_avg, mask = mask)
+
+def update_fn(
+        p: torch.Tensor,
+        grad: torch.Tensor,
+        exp_avg: torch.Tensor,
+        lr: float,
+        wd: float,
+        beta1: float,
+        beta2: float
+):
+    assert all([t.is_cuda for t in (p, grad, exp_avg)])
+    n_elements = p.numel()
+
+    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
+
+    update_fn_kernel[grid](
+        p,
+        grad,
+        exp_avg,
+        lr,
+        wd,
+        beta1,
+        beta2,
+        n_elements
+    )

src/train/opt/opt_modules/lookahead_opt.py

@@ -0,0 +1,153 @@
+#https://github.com/kozistr/pytorch_optimizer/blob/main/pytorch_optimizer/optimizer/lookahead.py
+
+
+
+
+import torch
+from torch.optim import Optimizer
+from collections import defaultdict
+from typing import Any, Dict
+from src.utils.utils_types import OptFloat, OptLossClosure, State
+
+from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Union
+
+from torch import Tensor
+
+Params = Union[Iterable[Tensor], Iterable[Dict[str, Any]]]
+
+from collections import defaultdict
+from typing import Callable, Dict
+
+import torch
+from src.train.opt.opt_modules.base_optimiser import BaseOptimizer, CLOSURE, DEFAULTS, LOSS, OPTIMIZER, STATE
+
+
+class Lookahead(BaseOptimizer):
+    r"""k steps forward, 1 step back.
+
+    :param optimizer: OPTIMIZER. base optimizer.
+    :param k: int. number of lookahead steps.
+    :param alpha: float. linear interpolation factor.
+    :param pullback_momentum: str. change to inner optimizer momentum on interpolation update.
+    """
+
+    def __init__(
+            self,
+            optimizer: OPTIMIZER,
+            k: int = 5,
+            alpha: float = 0.5,
+            pullback_momentum: str = 'none',
+            **kwargs,
+    ) -> None:
+        self.validate_positive(k, 'k')
+        self.validate_range(alpha, 'alpha', 0.0, 1.0)
+        self.validate_options(pullback_momentum, 'pullback_momentum', ['none', 'reset', 'pullback'])
+
+        self._optimizer_step_pre_hooks: Dict[int, Callable] = {}
+        self._optimizer_step_post_hooks: Dict[int, Callable] = {}
+
+        self.alpha = alpha
+        self.k = k
+        self.pullback_momentum = pullback_momentum
+
+        self.optimizer = optimizer
+
+        self.state: STATE = defaultdict(dict)
+
+        for group in self.param_groups:
+            if 'counter' not in group:
+                group['counter'] = 0
+
+            for p in group['params']:
+                state = self.state[p]
+                state['slow_params'] = torch.empty_like(p)
+                state['slow_params'].copy_(p)
+                if self.pullback_momentum == 'pullback':
+                    state['slow_momentum'] = torch.zeros_like(p)
+
+        self.defaults: DEFAULTS = {
+            'lookahead_alpha': alpha,
+            'lookahead_k': k,
+            'lookahead_pullback_momentum': pullback_momentum,
+            **optimizer.defaults,
+        }
+
+    @property
+    def param_groups(self):
+        return self.optimizer.param_groups
+
+    def __getstate__(self):
+        return {
+            'state': self.state,
+            'optimizer': self.optimizer,
+            'alpha': self.alpha,
+            'k': self.k,
+            'pullback_momentum': self.pullback_momentum,
+        }
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['counter'] = 0
+
+    def backup_and_load_cache(self):
+        r"""Backup cache parameters."""
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+                state['backup_params'] = torch.empty_like(p)
+                state['backup_params'].copy_(p)
+                p.data.copy_(state['slow_params'])
+
+    def clear_and_load_backup(self):
+        r"""Load backup parameters."""
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+                p.data.copy_(state['backup_params'])
+                del state['backup_params']
+
+    def state_dict(self) -> STATE:
+        return self.optimizer.state_dict()
+
+    def load_state_dict(self, state: STATE):
+        r"""Load state."""
+        self.optimizer.load_state_dict(state)
+
+    @torch.no_grad()
+    def zero_grad(self):
+        self.optimizer.zero_grad(set_to_none=True)
+
+    @torch.no_grad()
+    def update(self, group: Dict):
+        for p in group['params']:
+            if p.grad is None:
+                continue
+
+            state = self.state[p]
+
+            slow = state['slow_params']
+
+            p.mul_(self.alpha).add_(slow, alpha=1.0 - self.alpha)
+            slow.copy_(p)
+
+            if 'momentum_buffer' not in self.optimizer.state[p]:
+                self.optimizer.state[p]['momentum_buffer'] = torch.zeros_like(p)
+
+            if self.pullback_momentum == 'pullback':
+                internal_momentum = self.optimizer.state[p]['momentum_buffer']
+                self.optimizer.state[p]['momentum_buffer'] = internal_momentum.mul_(self.alpha).add_(
+                    state['slow_momentum'], alpha=1.0 - self.alpha
+                )
+                state['slow_momentum'] = self.optimizer.state[p]['momentum_buffer']
+            elif self.pullback_momentum == 'reset':
+                self.optimizer.state[p]['momentum_buffer'] = torch.zeros_like(p)
+
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = self.optimizer.step(closure)
+        for group in self.param_groups:
+            group['counter'] += 1
+            if group['counter'] >= self.k:
+                group['counter'] = 0
+                self.update(group)
+        return loss

src/train/opt/opt_modules/lookahead_opt_old.py

@@ -0,0 +1,130 @@
+#https://github.com/kozistr/pytorch_optimizer/blob/main/pytorch_optimizer/optimizer/lookahead.py
+
+
+
+
+import torch
+from torch.optim import Optimizer
+from collections import defaultdict
+from typing import Any, Dict
+from src.utils.utils_types import OptFloat, OptLossClosure, State
+
+from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Union
+
+from torch import Tensor
+
+Params = Union[Iterable[Tensor], Iterable[Dict[str, Any]]]
+
+
+class Lookahead(Optimizer):
+    r"""Implements Lookahead optimization algorithm.
+    It has been proposed in `Lookahead Optimizer: k steps forward, 1
+    step back`__
+    (https://arxiv.org/pdf/1907.08610.pdf)
+    Arguments:
+        optimizer: base inner optimizer optimize, like Yogi, DiffGrad or Adam.
+        k: number of lookahead steps (default: 5)
+        alpha: linear interpolation factor. 1.0 recovers the inner optimizer.
+            (default: 5)
+    Example:
+        >>> optimizer = optim.Lookahead(optimizer, k=5, alpha=0.5)
+        >>> optimizer.zero_grad()
+        >>> loss_fn(model(input), target).backward()
+        >>> optimizer.step()
+    __ https://arxiv.org/abs/1907.08610
+    Note:
+        Reference code: https://github.com/jettify/pytorch-optimizer/blob/master/torch_optimizer/lookahead.py
+    """
+
+    def __init__(
+            self, optimizer: Optimizer, k: int = 5, alpha: float = 0.5
+    ) -> None:
+        if k < 0.0:
+            raise ValueError('Invalid number of lookahead steps: {}'.format(k))
+        if alpha < 0:
+            raise ValueError(
+                'Invalid linear interpolation factor: {}'.format(alpha)
+            )
+
+        self.optimizer = optimizer
+        self.k = k
+        self.alpha = alpha
+        self.param_groups = self.optimizer.param_groups
+        self.state = defaultdict(dict)
+        self.fast_state = self.optimizer.state
+        for group in self.param_groups:
+            group['counter'] = 0
+        self.defaults = {'k': k, 'alpha': alpha, **optimizer.defaults}
+
+    def _update(self, group: Dict[str, Any]) -> None:
+        for fast in group['params']:
+            param_state = self.state[fast]
+            if 'slow_param' not in param_state:
+                param_state['slow_param'] = torch.clone(fast.data).detach()
+
+            slow = param_state['slow_param']
+            fast.data.mul_(self.alpha).add_(slow, alpha=1.0 - self.alpha)
+            slow.data.copy_(fast)
+
+    def step(self, closure: OptLossClosure = None) -> OptFloat:
+        r"""Performs a single optimization step.
+        Arguments:
+            closure: A closure that reevaluates the model and returns the loss.
+        """
+        loss = self.optimizer.step(closure=closure)
+        for group in self.param_groups:
+            if group['counter'] == 0:
+                self._update(group)
+            group['counter'] += 1
+            group['counter'] %= self.k
+        return loss
+
+    def state_dict(self) -> State:
+        r"""Returns the state of the optimizer as a :class:`dict`.
+        It contains two entries:
+        * state - a dict holding current optimization state. Its content
+            differs between optimizer classes.
+        * param_groups - a dict containing all parameter groups
+        """
+        slow_state_dict = super(Lookahead, self).state_dict()
+        fast_state_dict = self.optimizer.state_dict()
+        fast_state = fast_state_dict['state']
+        param_groups = fast_state_dict['param_groups']
+        return {
+            'fast_state': fast_state,
+            'slow_state': slow_state_dict['state'],
+            'param_groups': param_groups,
+        }
+
+    def load_state_dict(self, state_dict: State) -> None:
+        r"""Loads the optimizer state.
+        Arguments:
+            state_dict: optimizer state. Should be an object returned
+                from a call to :meth:`state_dict`.
+        """
+        slow_state_dict = {
+            'state': state_dict['slow_state'],
+            'param_groups': state_dict['param_groups'],
+        }
+        fast_state_dict = {
+            'state': state_dict['fast_state'],
+            'param_groups': state_dict['param_groups'],
+        }
+        super(Lookahead, self).load_state_dict(slow_state_dict)
+        self.optimizer.load_state_dict(fast_state_dict)
+        self.fast_state = self.optimizer.state
+
+    def zero_grad(self, set_to_none: bool = False) -> None:
+        r"""Clears the gradients of all optimized :class:`torch.Tensor` s."""
+        self.optimizer.zero_grad(set_to_none)
+
+    def __repr__(self) -> str:
+        base_str = self.optimizer.__repr__()
+        format_string = self.__class__.__name__ + ' ('
+        format_string += '\n'
+        format_string += 'k: {}\n'.format(self.k)
+        format_string += 'alpha: {}\n'.format(self.alpha)
+        format_string += base_str
+        format_string += '\n'
+        format_string += ')'
+        return format_string

src/train/opt/opt_modules/radam_opt.py

@@ -0,0 +1,119 @@
+import math
+import torch
+from torch.optim import Optimizer
+
+
+class RAdam(Optimizer):
+    r"""Implements RAdam algorithm.
+    It has been proposed in `ON THE VARIANCE OF THE ADAPTIVE LEARNING
+    RATE AND BEYOND(https://arxiv.org/pdf/1908.03265.pdf)`_.
+
+    Arguments:
+        params (iterable):      iterable of parameters to optimize or dicts defining
+                                parameter groups
+        lr (float, optional):   learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional):  coefficients used for computing
+                                                running averages of gradient and
+                                                its square (default: (0.9, 0.999))
+        eps (float, optional):  term added to the denominator to improve
+                                numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional):    whether to use the AMSGrad variant of this
+                                        algorithm from the paper `On the Convergence
+                                        of Adam and Beyond`_(default: False)
+
+        sma_thresh:             simple moving average threshold.
+                                Length till where the variance of adaptive lr is intracable.
+                                Default: 4 (as per paper)
+    """
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0, amsgrad=False, sma_thresh=4):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay, amsgrad=amsgrad)
+        super(RAdam, self).__init__(params, defaults)
+
+        self.radam_buffer = [[None, None, None] for ind in range(10)]
+        self.sma_thresh = sma_thresh
+
+    def __setstate__(self, state):
+        super(RAdam, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                # Perform optimization step
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
+                amsgrad = group['amsgrad']
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                if amsgrad:
+                    max_exp_avg_sq = state['max_exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+                old = p.data.float()
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+
+                buffer = self.radam_buffer[int(state['step']%10)]
+
+                if buffer[0] == state['step']:
+                    sma_t, step_size = buffer[1], buffer[2]
+                else:
+                    sma_max_len = 2/(1-beta2) - 1
+                    beta2_t = beta2 ** state['step']
+                    sma_t = sma_max_len - 2 * state['step'] * beta2_t /(1 - beta2_t)
+                    buffer[0] = state['step']
+                    buffer[1] = sma_t
+
+                    if sma_t > self.sma_thresh :
+                        rt = math.sqrt(((sma_t - 4) * (sma_t - 2) * sma_max_len)/((sma_max_len -4) * (sma_max_len - 2) * sma_t))
+                        step_size = group['lr'] * rt * math.sqrt((1 - beta2_t)) / (1 - beta1 ** state['step'])
+                    else:
+                        step_size = group['lr'] / (1 - beta1 ** state['step'])
+                    buffer[2] = step_size
+
+                if group['weight_decay'] != 0:
+                    p.data.add_(-group['weight_decay'] * group['lr'], old)
+
+                if sma_t > self.sma_thresh :
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+                    p.data.addcdiv_(-step_size, exp_avg, denom)
+                else:
+                    p.data.add_(-step_size, exp_avg)
+
+        return

src/train/opt/opt_modules/weight_decay.py

@@ -0,0 +1,65 @@
+from typing import Optional, List, Tuple
+import torch
+
+
+
+def set_weight_decay(
+        model: torch.nn.Module,
+        weight_decay: float,
+        norm_weight_decay: Optional[float] = None,
+        norm_classes: Optional[List[type]] = None,
+        custom_keys_weight_decay: Optional[List[Tuple[str, float]]] = None,
+):
+    if not norm_classes:
+        norm_classes = [
+            torch.nn.modules.batchnorm._BatchNorm,
+            torch.nn.LayerNorm,
+            torch.nn.GroupNorm,
+            torch.nn.modules.instancenorm._InstanceNorm,
+            torch.nn.LocalResponseNorm,
+        ]
+    norm_classes = tuple(norm_classes)
+
+    params = {
+        "other": [],
+        "norm": [],
+    }
+    params_weight_decay = {
+        "other": weight_decay,
+        "norm": norm_weight_decay,
+    }
+    custom_keys = []
+    if custom_keys_weight_decay is not None:
+        for key, weight_decay in custom_keys_weight_decay:
+            params[key] = []
+            params_weight_decay[key] = weight_decay
+            custom_keys.append(key)
+
+    def _add_params(module, prefix=""):
+        for name, p in module.named_parameters(recurse=False):
+            if not p.requires_grad:
+                continue
+            is_custom_key = False
+            for key in custom_keys:
+                target_name = f"{prefix}.{name}" if prefix != "" and "." in key else name
+                if key == target_name:
+                    params[key].append(p)
+                    is_custom_key = True
+                    break
+            if not is_custom_key:
+                if norm_weight_decay is not None and isinstance(module, norm_classes):
+                    params["norm"].append(p)
+                else:
+                    params["other"].append(p)
+
+        for child_name, child_module in module.named_children():
+            child_prefix = f"{prefix}.{child_name}" if prefix != "" else child_name
+            _add_params(child_module, prefix=child_prefix)
+
+    _add_params(model)
+
+    param_groups = []
+    for key in params:
+        if len(params[key]) > 0:
+            param_groups.append({"params": params[key], "weight_decay": params_weight_decay[key]})
+    return param_groups

src/train/train_naruto.py

@@ -0,0 +1,345 @@
+# https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html
+# https://www.learnpytorch.io/06_pytorch_transfer_learning/
+
+from __future__ import print_function
+from __future__ import division
+
+import torch
+import torch.nn as nn
+import numpy as np
+import gin
+from pathlib import Path
+import matplotlib.pyplot as plt
+from torch.utils.data import DataLoader
+import torch.backends.cudnn as cudnn
+import os
+###
+from src.conf.train_conf import TrainConfig
+from src.ds.naruto_dl import ObjectNarutoDataset
+
+
+from src.utils.utils_train import enable_tf32, clear_memory, get_device
+from src.utils.utils_file_dir import create_dir_tree
+from src.train.criterion.metrics.accuracy import accuracy as accuracy_func
+from src.train.criterion import get_loss_func
+from src.train.opt import create_optimizer
+
+from src.models.bb.load_models import load_models as load_models_func
+from src.models.bb.make_fe import make_fe_from_classification_model
+from src.utils.utils_log import create_logger, double_dash_line
+from torchmetrics import F1Score
+
+cudnn.benchmark = True
+plt.ion()   # interactive mode
+num_cores = 6 #os.cpu_count()
+
+class TrainerNaruto():
+    """
+    Args:
+       model_name: str: name of current model
+       albert_config: AlbertConf : albert config object
+       train_config: TrainConf: training config object
+    """
+    def __init__(self, model_name, model, train_config, input_config, fdir_config, device):
+        self.config = train_config  # training config
+        self.model_name = model_name
+        self.model = model
+        self.device = device
+        self.show_per_step = self.config.show_per_step
+        self.logger_name = f'{type(self).__name__}_{self.model_name}'
+        self.logger = create_logger(log_name=self.logger_name)
+        self.input_size = self.config.input_size  # vocabulary size
+        # path to model directory to hold output results
+        model_dir =  f'{self.config.out_dir}/{model_name}_{self.input_size}'
+        # path to store info of model
+        self.meta_dir = f'{model_dir}/meta'
+        # path to store tensotboards results
+        self.tb_dir = f'{self.meta_dir}/tb'
+        # weights directory
+        self.weights_dir = f'{model_dir}/weights'
+        #
+        # self.meta_test = f'{self.meta_dir}/test.csv'
+        # self.meta_train = f'{self.meta_dir}/train.csv'
+        #
+        # arr dirs to hold results
+        arr_dirs = [model_dir, self.meta_dir, self.tb_dir, self.weights_dir]
+        self.logger.info('Creating subdirs for train process...')
+        # create dir trees if absent
+        create_dir_tree(arr_dirs)
+        #
+        #!todo create scaler
+        # norm value for gradient clipping op
+        self.max_grad_norm = self.config.max_grad_norm
+        # num steps for gradient clipping
+        self.gradient_accumulator_steps = self.config.gradient_accumulation_steps
+        # compute global batches size
+        self.global_train_batch_size = self.config.train_batch_size * self.gradient_accumulator_steps
+        self.logger.info("Initializing Datasets and Dataloaders...")
+        # Create training and validation dataloaders
+        ds_train = ObjectNarutoDataset(mode='train', fdir_cfg=fdir_config, input_cfg=input_config)
+        ds_test = ObjectNarutoDataset(mode='test', fdir_cfg=fdir_config, input_cfg=input_config)
+        # num classes
+        self.num_classes = ds_train.lab2id.classes_.size
+        self.config.data_size = len(ds_train)
+        self.train_dl = DataLoader(ds_train, batch_size=self.global_train_batch_size, shuffle=True,
+                                   num_workers=num_cores, pin_memory=True, #drop_last=True,
+                                   persistent_workers=False)
+        self.test_dl = DataLoader(ds_test, batch_size=self.config.test_batch_size, shuffle=False,
+                                  num_workers=num_cores, pin_memory=True, persistent_workers=False)
+        self.logger.info("Initialized Datasets and Dataloaders!")
+        # initialize number steps for training procedure
+        self.num_epochs = self.config.num_epochs
+        self.num_batches_ds_train = len(self.train_dl)
+        self.num_batches_ds_test = len(self.test_dl)
+        self.num_batches_per_epoch = self.num_batches_ds_train // self.gradient_accumulator_steps
+        self.num_batches_total = self.num_batches_per_epoch * self.num_epochs
+        # number of warmup steps for linear loss
+        self.num_warmup_steps = self.config.num_warmup_steps
+        if self.num_warmup_steps == 0:
+            if self.config.warmup_proportion == 0:
+                self.config.warmup_proportion = 0.1
+            self.num_warmup_steps = round(self.num_batches_total * self.config.warmup_proportion)
+        #
+        self.logger.info('Format loaded model for new classification task...')
+        # loss_func
+        self.loss_func = get_loss_func(loss_type=self.config.loss_func_type, eps=self.config.opt_eps,
+                                       label_smooth_coef=self.config.label_smooth_coef,)
+
+        self.f1_micro_score_func = F1Score(task="multiclass", num_classes=self.num_classes, average='micro').to(self.device)
+        self.f1_macro_score_func = F1Score(task="multiclass", num_classes=self.num_classes, average='macro').to(self.device)
+        self.f1_weighted_score_func = F1Score(task="multiclass", num_classes=self.num_classes, average='weighted').to(self.device)
+
+
+        self.model = make_fe_from_classification_model(model_ft=model, model_name=model_name,
+                                                       num_class=self.num_classes, drop_rate=train_config.drop_rate,
+                                                       )
+        self.model = self.model.to(self.device)
+
+        self.logger.info('Creating optimizers and scheduler for learning rate')
+        # optimizer
+        self.opt_func, self.lr_scheduler = create_optimizer(model=self.model, warmup_steps=self.num_warmup_steps,
+                                                            initial_lr=self.config.initial_lr, max_lr=self.config.max_lr,
+                                                            weight_decay_rate=self.config.weight_decay_rate,
+                                                            lr_decay_rate=self.config.lr_decay_rate,
+                                                            opt_beta1=self.config.opt_beta1, opt_beta2=self.config.opt_beta2,
+                                                            opt_eps=self.config.opt_eps,
+                                                            opt_type=self.config.opt_type, lr_schedule_type=self.config.lr_schedule_type,
+                                                            momentum=self.config.momentum, use_lookahead_opt=self.config.use_lookahead_opt)
+
+
+        self.logger.info(f'\n'
+                         f'Model_name: {model_name}\n'
+                         f'Num epochs: {self.config.num_epochs}\n'
+                         f'Train batch size: {self.config.train_batch_size}\n'
+                         f'Global train batch size: {self.global_train_batch_size}\n'
+                         f'Number of batches in train dataset: {self.num_batches_ds_train}\n'
+                         f'Number of batches in one epoch: {self.num_batches_per_epoch}\n'
+                         f'Total number of batches for {self.num_epochs} epochs: {self.num_batches_total}\n'
+                         f'Num warmup steps: {self.num_warmup_steps}\n'
+                         f'Optimizer: {self.config.opt_type}\n'
+                         f'Scheduler learning rate: {self.config.lr_schedule_type}\n'
+                         f'Loss function: {self.config.loss_func_type}\n'
+                         f'Num classes: {self.num_classes}\n\n'
+                         )
+
+    def train_model(self):
+        self.best_val_loss = 99999
+        self.count_batches = 0
+        self.model = self.model.to(self.device)
+        for epoch in range(self.num_epochs):
+            self.cur_epoch = epoch + 1
+            self.logger.info(f'EPOCH: {self.cur_epoch} of {self.num_epochs}')
+            self.logger_epoch = create_logger(f'{self.logger_name}_train_epoch_{self.cur_epoch}')
+            self.model.train(True)
+            avg_loss, avg_acc1, avg_f1_micro, avg_f1_macro, avg_f1_weighted = self.train_epoch_step()
+            self.logger_epoch.info(double_dash_line())
+            self.logger_epoch.info(f"Avg_loss: {avg_loss}\nAvg_acc1:{avg_acc1}\n"
+                                   f"Avg_f1_micro: {avg_f1_micro}\tAvg_f1_macro: {avg_f1_macro}\tAvg_f1_w: {avg_f1_weighted}\t")
+            self.logger_epoch.info(double_dash_line())
+
+            with open(f'{self.meta_dir}/meta_{self.model_name}.txt', 'a') as o_f:
+                o_f.write(f'{self.cur_epoch}\t{avg_loss}\t{avg_acc1}\t{avg_f1_micro}\t{avg_f1_macro}\t{avg_f1_weighted}\n')
+
+            val_loss = self.eval_model()
+            if val_loss < self.best_val_loss:
+                self.best_val_loss = val_loss
+                self.logger_epoch.info('Saving model ...')
+                state = {
+                    'epoch': epoch,
+                    'state_dict': self.model.state_dict(),
+                    'optimizer': self.opt_func.state_dict(),
+                    'lr_scheduler' : self.lr_scheduler.state_dict(),
+                }
+                torch.save(state, f'{self.weights_dir}/best_model.pth')
+
+            lr = self.lr_scheduler.get_lr()
+            self.logger_epoch.info(f'Learning rate: {lr}')
+            with open(f'{self.meta_dir}/meta_lr_{self.model_name}.txt', 'a') as o_f:
+                o_f.write(f'{self.cur_epoch}\t{lr}\n')
+
+
+            self.lr_scheduler.step(val_loss)
+
+    def train_epoch_step(self):
+        running_loss = 0.
+        running_acc1 = 0.
+
+        running_f1_micro = 0.
+        running_f1_macro = 0.
+        running_f1_weighted = 0.
+
+        last_loss = 0.
+        last_acc1 = 0.
+
+        last_f1_micro = 0.
+        last_f1_macro = 0.
+        last_f1_weighted = 0.
+
+
+        epoch_loss = 0.
+        epoch_acc1 = 0.
+
+        epoch_f1_micro = 0.
+        epoch_f1_macro = 0.
+        epoch_f1_weighted = 0.
+
+        for i_batch, batch in enumerate(self.train_dl):
+            self.count_batches += 1
+
+            image, label = batch
+            img_tensor = image.to(self.device)
+            lab_tensor = label.to(self.device)
+
+            # Make predictions for this batch
+            res = self.model(img_tensor)
+            # Compute the loss and its gradients
+            loss = self.loss_func(res, lab_tensor)
+            self.opt_func.zero_grad()
+            loss.backward()
+            # Adjust learning weights
+            self.opt_func.step()
+            acc1 = accuracy_func(res, lab_tensor, topk=(1, ))[0]
+            f1_micro_val = self.f1_micro_score_func(res, lab_tensor)
+            f1_macro_val = self.f1_macro_score_func(res, lab_tensor)
+            f1_weighted_val = self.f1_weighted_score_func(res, lab_tensor)
+
+            running_f1_micro += f1_micro_val
+            running_f1_macro += f1_macro_val
+            running_f1_weighted += f1_weighted_val
+
+            running_loss += loss.item()
+            running_acc1 += acc1
+
+
+            epoch_loss += loss.item()
+            epoch_acc1 += acc1
+
+            epoch_f1_micro += f1_micro_val
+            epoch_f1_macro += f1_macro_val
+            epoch_f1_weighted += f1_weighted_val
+
+            if (i_batch+1) % self.show_per_step == 0:
+                last_loss = running_loss / self.show_per_step # loss per batch
+                last_acc1 = running_acc1 / self.show_per_step # loss per batch
+
+                last_f1_macro = running_f1_macro / self.show_per_step
+                last_f1_micro = running_f1_micro / self.show_per_step
+                last_f1_weighted = running_f1_weighted / self.show_per_step
+
+
+                self.logger_epoch.info(f'batch_{i_batch+1} of {self.num_batches_per_epoch}\n'
+                                       f'loss: {last_loss}\n'
+                                       f'acc1:{last_acc1}\n'
+                                       f'f1_micro: {last_f1_micro} | f1_macro: {last_f1_macro} | f1_weighted:{last_f1_weighted}'
+                                       )
+                with open(f'{self.meta_dir}/meta_{self.model_name}_epoch_{self.cur_epoch}.txt', 'a') as o_f:
+                    o_f.write(f'{i_batch+1}\t{last_loss}\t{last_acc1}\t{last_f1_micro}\t{last_f1_macro}\t{last_f1_weighted}\n')
+
+                running_loss = 0.
+                running_f1_weighted = 0.
+                running_f1_macro = 0.
+                running_f1_micro = 0.
+                running_acc1 = 0.
+
+
+            if self.count_batches < self.num_warmup_steps:
+                self.lr_scheduler.step()
+
+        epoch_loss /= self.num_batches_per_epoch
+        epoch_f1_micro /= self.num_batches_per_epoch
+        epoch_f1_macro /= self.num_batches_per_epoch
+        epoch_f1_weighted /= self.num_batches_per_epoch
+        epoch_acc1 /= self.num_batches_per_epoch
+
+        return epoch_loss, epoch_acc1, epoch_f1_micro, epoch_f1_macro, epoch_f1_weighted
+
+    def eval_model(self):
+        self.model.eval()
+
+        val_loss = 0.
+        val_acc1 = 0.
+        val_f1_micro = 0.
+        val_f1_macro = 0.
+        val_f1_weighted = 0.
+
+        with torch.no_grad():
+            for i_batch, batch in enumerate(self.test_dl):
+                image, label = batch
+                img_tensor = image.to(self.device)
+                lab_tensor = label.to(self.device)
+                # Make predictions for this batch
+                res = self.model(img_tensor)
+                # Compute the loss and its gradients
+                loss = self.loss_func(res, lab_tensor)
+                acc1 = accuracy_func(res, lab_tensor, topk=(1,))[0]
+                f1_micro = self.f1_micro_score_func(res, lab_tensor)
+                f1_macro = self.f1_macro_score_func(res, lab_tensor)
+                f1_weighted = self.f1_weighted_score_func(res, lab_tensor)
+                val_loss += loss.item()
+                val_acc1 += acc1
+                val_f1_micro += f1_micro
+                val_f1_macro += f1_macro
+                val_f1_weighted += f1_weighted
+
+
+        val_loss /= self.num_batches_ds_test
+        val_acc1 /= self.num_batches_ds_test
+
+        val_f1_micro /= self.num_batches_ds_test
+        val_f1_macro /= self.num_batches_ds_test
+        val_f1_weighted /= self.num_batches_ds_test
+
+        self.logger_epoch.info(f'Test loss: {val_loss}\nacc1:{val_acc1}\n'
+                               f'f1_micro: {val_f1_micro}\tf1_macro: {val_f1_macro}\tf1_w: {val_f1_weighted}\n')
+        with open(f'{self.meta_dir}/meta_{self.model_name}_test.txt', 'a') as o_f:
+            o_f.write(f'{self.cur_epoch}\t{val_loss}\t{val_acc1}\t{val_f1_micro}\t{val_f1_macro}\t{val_f1_weighted}\n')
+        return val_loss
+
+def main_func():
+    cudnn.benchmark = True
+    logger = create_logger("InitOp")
+    logger.info(double_dash_line())
+    from src.conf.base_conf import get_fdir_cfg, get_inputval_cfg
+    from src.conf.train_conf import get_train_cfg
+    from src.conf.nn_conf import get_nn_cfg
+    from src.utils.utils_file_dir import get_proj_dir
+    proj_dir = get_proj_dir()
+    path2cfg = fr'{get_proj_dir()}in/config_files/'
+    fdir_conf = get_fdir_cfg(path2cfg)
+    input_conf = get_inputval_cfg(path2cfg)
+    nn_conf = get_nn_cfg(path2cfg)
+    train_conf = get_train_cfg(path2cfg)
+    logger.info('Loaded configs!')
+    clear_memory()
+    logger.info('Cleared memory')
+    device = get_device()
+    enable_tf32()
+    logger.info(f'Loading model --> {nn_conf.model_name}...')
+    model = load_models_func(nn_conf)#.to(device)
+    logger.info(f'Model {nn_conf.model_name} loaded!')
+    trainer_obj = TrainerNaruto(model_name=nn_conf.model_name, model=model,
+                             train_config=train_conf, input_config=input_conf, fdir_config=fdir_conf, device=device)
+    trainer_obj.train_model()
+
+if __name__ == "__main__":
+    main_func()

src/utils/utils_file_dir.py

@@ -0,0 +1,151 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from pathlib import Path
+import sys, os
+from natsort import natsorted
+from sklearn.model_selection import train_test_split
+import pandas as pd
+import shutil
+
+
+def get_subdirs(folder):
+    return [ f.path for f in os.scandir(folder) if f.is_dir() ]
+
+def del_empty_folders(fpath):
+    """
+    del empty folder
+    :param fpath:
+    :return:
+    """
+    for fpath in get_subdirs(fpath):
+        if len(os.listdir(fpath)) == 0: # Check if the folder is empty
+            print(f'Remove {fpath}')
+            shutil.rmtree(fpath) # If so, delete it
+
+def convert_size_to_mb(size_bytes):
+    """"""
+    if size_bytes == 0:
+        return 0
+    else:
+        return size_bytes / 1e6
+
+def get_size_bytes(arr):
+    """"""
+    return sys.getsizeof(arr)
+
+
+def check_file_exist(path):
+    return os.path.isfile(path)
+
+def get_files(path_dir, format=('wav')):
+    """
+    function tp get array of files with some format
+    Args:
+	    path_dir: str: path to directory
+	    format: format of file
+	Returns:
+	    arr: sorted array of files
+	"""
+    arr = []
+    p_opf = Path(path_dir)
+    if (p_opf.is_dir() == True and len(os.listdir(path_dir)) > 0):
+        for root, dirs, files in os.walk(path_dir):
+            for i in files:
+                if i.endswith(format):
+                    arr.append(root + '/' + i)
+    return natsorted(arr)
+
+def create_dir_tree(arr_dirs):
+    # create dir trees if absent
+    for dir in arr_dirs:
+        if not Path(dir).is_dir():
+            Path(dir).mkdir(parents=True, exist_ok=True)
+
+def get_proj_dir(name_proj=None):
+    """
+    func which allow to get initial dir of project by name
+    name_proj : name of project
+    :return:
+    """
+    '''dir_ = os.path.dirname(os.path.realpath(__file__))
+
+    split_dir = dir_.split('/')
+
+    k = 0
+    for i, d in enumerate(split_dir):
+        if d == name_proj:
+            k = i
+
+    dir_name = '/'.join(split_dir[:k + 1]) + '/'
+    return dir_name'''
+    return os.path.dirname(os.path.abspath(__file__)).split('src')[0]
+
+
+
+def split_stratified_into_train_val_test(df_input, stratify_colname='y',
+                                         frac_train=0.8, frac_val=0.05, frac_test=0.15,
+                                         random_state=None):
+    '''
+    Splits a Pandas dataframe into three subsets (train, val, and test)
+    following fractional ratios provided by the user, where each subset is
+    stratified by the values in a specific column (that is, each subset has
+    the same relative frequency of the values in the column). It performs this
+    splitting by running train_test_split() twice.
+
+    Parameters
+    ----------
+    df_input : Pandas dataframe
+        Input dataframe to be split.
+    stratify_colname : str
+        The name of the column that will be used for stratification. Usually
+        this column would be for the label.
+    frac_train : float
+    frac_val   : float
+    frac_test  : float
+        The ratios with which the dataframe will be split into train, val, and
+        test data. The values should be expressed as float fractions and should
+        sum to 1.0.
+    random_state : int, None, or RandomStateInstance
+        Value to be passed to train_test_split().
+
+    Returns
+    -------
+    df_train, df_val, df_test :
+        Dataframes containing the three splits.
+    '''
+
+    if frac_train + frac_val + frac_test != 1.0:
+        raise ValueError('fractions %f, %f, %f do not add up to 1.0' % \
+                         (frac_train, frac_val, frac_test))
+
+    if stratify_colname not in df_input.columns:
+        raise ValueError('%s is not a column in the dataframe' % (stratify_colname))
+
+    X = df_input # Contains all columns.
+    y = df_input[[stratify_colname]] # Dataframe of just the column on which to stratify.
+    # Split original dataframe into train and temp dataframes.
+    df_train, df_temp, y_train, y_temp = train_test_split(X,
+                                                          y,
+                                                          stratify=y,
+                                                          test_size=(1.0 - frac_train),
+                                                          random_state=random_state)
+
+    # Split the temp dataframe into val and test dataframes.
+    relative_frac_test = frac_test / (frac_val + frac_test)
+    df_val, df_test, y_val, y_test = train_test_split(df_temp,
+                                                      y_temp,
+                                                      stratify=y_temp,
+                                                      test_size=relative_frac_test,
+                                                      random_state=random_state)
+
+    assert len(df_input) == len(df_train) + len(df_val) + len(df_test)
+
+    return df_train, df_val, df_test
+
+if __name__ == "__main__":
+    pass

src/utils/utils_img.py

@@ -0,0 +1,172 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+
+import torch
+from torchvision.utils import draw_bounding_boxes
+from PIL import Image
+from PIL import ImageOps
+import torchvision
+import cv2
+import matplotlib.pyplot as plt
+import albumentations as A
+import numpy as np
+
+IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm']
+
+'''
+# Get a batch of training data
+inputs, classes = next(iter(dataloaders['train']))
+
+# Make a grid from batch
+out = torchvision.utils.make_grid(inputs)
+
+imshow(out, title=[class_names[x] for x in classes])
+'''
+
+
+
+def id2name_img_coco(img_id):
+    # 000000055022
+    return f"{img_id:012d}"
+
+def save_img_cv2(path, img):
+    cv2.imwrite(path, img)
+
+def crop_bbox_224(fimg, labels, bboxes):
+    transform = A.Compose([
+        A.Resize(height=255, width=255),
+        A.CenterCrop(height=224, width=224),
+    ], bbox_params=A.BboxParams(format='coco', label_fields=['class_labels']))
+
+    image = cv2.imread(fimg)
+    #image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    # visualize_bbox_plt(image, bboxes, labels)
+    transformed = transform(image=image, bboxes=bboxes, class_labels=labels)
+    transformed_image = transformed['image']
+    transformed_bboxes = transformed['bboxes']
+    transformed_class_labels = transformed['class_labels']
+    #visualize_bbox_plt(transformed_image, transformed_bboxes, transformed_class_labels)
+    return transformed_image, transformed_bboxes, transformed_class_labels
+
+
+
+def is_image_file(filename):
+    """Checks if a file is an image.
+    Args:
+        filename (string): path to a file
+    Returns:
+        bool: True if the filename ends with a known image extension
+    """
+    filename_lower = filename.lower()
+    return any(filename_lower.endswith(ext) for ext in IMG_EXTENSIONS)
+
+def convert_bboxes():
+    pass
+
+def pil_loader(path):
+    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
+    with open(path, 'rb') as f:
+        img = Image.open(f)
+        return img.convert('RGB')
+
+
+def crop_by_bbox_with_fill(h=224, w=224, img_path='', bbox=[], color='black', fill=False):
+    # pascal-voc format
+    x_min, y_min, x_max, y_max = bbox
+    img = Image.open(img_path)
+    img = img.convert('RGB')
+    # width, height = img.size
+    # width = x_max - x_min
+    # height = y_max - y_min
+    img = img.crop([x_min, y_min, x_max, y_max])
+    if fill:
+        try:
+            img = ImageOps.pad(img, (h, w), color=color)
+        except:
+            img = img.resize((h, w))
+    return img
+
+def draw_and_show_bbox_tv(img_path='', bbox='', flag_show=True):
+    """
+    :param img_path:
+    :param bbox:
+    :return:
+    """
+    img = Image.open(img_path)
+    img = img.convert('RGB')
+    img = torchvision.transforms.PILToTensor()(img)
+    box = torch.tensor(bbox)
+    box = box.unsqueeze(0)
+    # draw bounding box and fill color
+    img = draw_bounding_boxes(img, box, width=5, colors="green", fill=True)
+    # transform this image to PIL image
+    img = torchvision.transforms.ToPILImage()(img)
+    # display output
+    img.show()
+    return img
+
+def imshow(inp, title=None):
+    """Imshow for Tensor."""
+    inp = inp.numpy().transpose((1, 2, 0))
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
+    inp = std * inp + mean
+    inp = np.clip(inp, 0, 1)
+    plt.imshow(inp)
+    if title is not None:
+        plt.title(title)
+    plt.pause(0.001)  # pause a bit so that plots are updated
+
+
+def visualize_bbox_func(img, bbox, class_name,  thickness=2):
+    """Visualizes a single bounding box on the image"""
+    # https://www.rapidtables.com/web/color/RGB_Color.html
+    BOX_COLOR = (228, 28, 28) # Red
+    TEXT_COLOR = (255, 255, 255) # White
+    x_min, y_min, w, h = bbox
+    x_min, x_max, y_min, y_max = int(x_min), int(x_min + w), int(y_min), int(y_min + h)
+    cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=BOX_COLOR, thickness=thickness)
+    ((text_width, text_height), _) = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.35, 1)
+    cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)), (x_min + text_width, y_min), BOX_COLOR, -1)
+    cv2.putText(
+        img,
+        text=class_name,
+        org=(x_min, y_min - int(0.3 * text_height)),
+        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+        fontScale=0.35,
+        color=TEXT_COLOR,
+        lineType=cv2.LINE_AA,
+    )
+    return img
+
+
+def visualize_bbox_plt(image, bboxes, labels):
+    img = image.copy()
+    for bbox, label in zip(bboxes, labels):
+        #class_name = category_id_to_name[category_id]
+        img = visualize_bbox_func(img, bbox, label)
+    plt.figure(figsize=(12, 12))
+    plt.axis('off')
+    plt.imshow(img)
+    #cv2.imwrite('1.png', img)
+
+def tensor_from_img(path):
+
+    # check for rgb mode
+    img = Image.open(path)
+    img = img.convert('RGB')
+    from torchvision import transforms
+    resize = transforms.Resize([224, 224])
+    to_tensor = transforms.ToTensor()
+    img = resize(img)
+    tensor = to_tensor(img)
+    tensor = tensor.unsqueeze(0)
+    return tensor
+
+if __name__ == "_main__":
+    pass

src/utils/utils_log.py

@@ -0,0 +1,84 @@
+import coloredlogs, logging
+import logging
+import time
+from typing import Optional
+import tqdm
+
+class TqdmLoggingHandler(logging.Handler):
+    def __init__(self, level=logging.NOTSET):
+        super().__init__(level)
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            tqdm.tqdm.write(msg)
+            self.flush()
+        except Exception:
+            self.handleError(record)
+
+def create_logger(log_name='default'):
+    # Added: Create logger and assign handler
+    logger = logging.getLogger(log_name)
+    color_log(logger)
+    #logger.addHandler(TqdmLoggingHandler())
+    return logger
+
+def color_log(logger):
+    coloredlogs.install(level='DEBUG')
+    coloredlogs.install(level='DEBUG', logger=logger)
+    coloredlogs.install(fmt='%(asctime)s,%(msecs)03d %(hostname)s %(name)s[%(process)d] %(levelname)s %(message)s')
+
+text_colors = {
+    'logs': '\033[34m',  # 033 is the escape code and 34 is the color code
+    'info': '\033[32m',
+    'warning': '\033[33m',
+    'error': '\033[31m',
+    'bold': '\033[1m',
+    'end_color': '\033[0m',
+    'light_red': '\033[36m'
+}
+
+def get_curr_time_stamp() -> str:
+    return time.strftime("%Y-%m-%d %H:%M:%S")
+
+def error(message: str) -> None:
+    time_stamp = get_curr_time_stamp()
+    error_str = text_colors['error'] + text_colors['bold'] + 'ERROR  ' + text_colors['end_color']
+    print('{} - {} - {}'.format(time_stamp, error_str, message))
+    print('{} - {} - {}'.format(time_stamp, error_str, 'Exiting!!!'))
+    exit(-1)
+
+def color_text(in_text: str) -> str:
+    return text_colors['light_red'] + in_text + text_colors['end_color']
+
+def log(message: str) -> None:
+    time_stamp = get_curr_time_stamp()
+    log_str = text_colors['logs'] + text_colors['bold'] + 'LOGS   ' + text_colors['end_color']
+    print('{} - {} - {}'.format(time_stamp, log_str, message))
+
+def warning(message: str) -> None:
+    time_stamp = get_curr_time_stamp()
+    warn_str = text_colors['warning'] + text_colors['bold'] + 'WARNING' + text_colors['end_color']
+    print('{} - {} - {}'.format(time_stamp, warn_str, message))
+
+def info(message: str, print_line: Optional[bool] = False) -> None:
+    time_stamp = get_curr_time_stamp()
+    info_str = text_colors['info'] + text_colors['bold'] + 'INFO   ' + text_colors['end_color']
+    print('{} - {} - {}'.format(time_stamp, info_str, message))
+    if print_line:
+        double_dash_line(dashes=150)
+
+def double_dash_line(dashes: Optional[int] = 75) -> None:
+    strok = text_colors['error'] + '=' * dashes + text_colors['end_color']
+    return strok
+
+def singe_dash_line(dashes: Optional[int] = 67) -> None:
+    print('-' * dashes)
+
+def print_header(header: str) -> None:
+    double_dash_line()
+    print(text_colors['info'] + text_colors['bold'] + '=' * 50 + str(header) + text_colors['end_color'])
+    double_dash_line()
+
+def print_header_minor(header: str) -> None:
+    print(text_colors['warning'] + text_colors['bold'] + '=' * 25 + str(header) + text_colors['end_color'])

src/utils/utils_nn.py

@@ -0,0 +1,10 @@
+
+
+
+def autopad(k, p=None, d=1):  # kernel, padding, dilation
+    """Pad to 'same' shape outputs."""
+    if d > 1:
+        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
+    if p is None:
+        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
+    return p

src/utils/utils_train.py

@@ -0,0 +1,40 @@
+import random
+import gc
+import sys
+import ctypes
+import numpy as np
+import torch
+
+def random_seed(seed=42, rank=0):
+    """
+    :param seed:
+    :param rank:
+    :return:
+    """
+    torch.manual_seed(seed + rank)
+    np.random.seed(seed + rank)
+    random.seed(seed + rank)
+
+def enable_tf32():
+    # This enables tf32 on Ampere GPUs which is only 8% slower than float16 and almost as accurate as float32
+    # This was a default in pytorch until 1.12
+    if torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.deterministic = False
+
+def clear_memory():
+    # clear memory
+    gc.collect()
+    torch.cuda.empty_cache()
+    if sys.platform == "linux" or sys.platform == "linux2":
+        # linux
+        libc = ctypes.CDLL("libc.so.6")
+        libc.malloc_trim(0)
+
+def get_device():
+    return torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+
+if __name__ == "__main__":
+    pass

src/utils/utils_types.py

@@ -0,0 +1,32 @@
+from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Union, List
+from types import FunctionType
+from collections import OrderedDict
+from torch import Tensor
+from itertools import repeat
+import collections.abc
+
+
+Params = Union[Iterable[Tensor], Iterable[Dict[str, Any]]]
+
+LossClosure = Callable[[], float]
+OptLossClosure = Optional[LossClosure]
+Betas2 = Tuple[float, float]
+State = Dict[str, Any]
+OptFloat = Optional[float]
+Nus2 = Tuple[float, float]
+
+_int_tuple_2_t = Union[int, Tuple[int, int]]
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
+    return parse
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple