From 5aaad307954f12c2b804ca88410131982a2b1625 Mon Sep 17 00:00:00 2001 From: whatispain Date: Thu, 2 Jul 2026 10:43:34 +0300 Subject: [PATCH] =?UTF-8?q?=D0=9F=D0=B5=D1=80=D0=B2=D1=8B=D0=B9=20=D0=BA?= =?UTF-8?q?=D0=BE=D0=BC=D0=BC=D0=B8=D1=82:=20=D0=BA=D0=BE=D0=B4=20=D0=BA?= =?UTF-8?q?=D0=BB=D0=B0=D1=81=D1=81=D0=B8=D1=84=D0=B8=D0=BA=D0=B0=D1=82?= =?UTF-8?q?=D0=BE=D1=80=D0=B0=20=D0=BD=D0=B0=D1=80=D1=83=D1=87=D0=BD=D1=8B?= =?UTF-8?q?=D1=85=20=D0=B7=D0=BD=D0=B0=D0=BA=D0=BE=D0=B2=20(Naruto)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 24 ++ env.yaml | 378 +++++++++++++++++ help | 5 + in/config_files/file_dir.gin | 3 + in/config_files/input.gin | 10 + in/config_files/nn.gin | 14 + in/config_files/train.gin | 57 +++ src/__init__.py | 0 src/conf/__init__.py | 0 src/conf/base_conf.py | 79 ++++ src/conf/nn_conf.py | 70 +++ src/conf/train_conf.py | 98 +++++ src/ds/__init__.py | 0 src/ds/gtauav_dl.py | 263 ++++++++++++ src/ds/naruto_dl.py | 257 +++++++++++ src/models/__init__.py | 0 src/models/bb/__init__.py | 0 src/models/bb/coca/__init__.py | 0 src/models/bb/coca/coca_model.py | 28 ++ src/models/bb/edgenext/edgenext_blocks.py | 335 +++++++++++++++ src/models/bb/edgenext/edgenext_model.py | 401 ++++++++++++++++++ src/models/bb/inception_next/__init__.py | 0 .../bb/inception_next/inception_next_model.py | 323 ++++++++++++++ src/models/bb/load_models.py | 88 ++++ src/models/bb/make_fe.py | 136 ++++++ src/train/criterion/__init__.py | 1 + src/train/criterion/init_loss.py | 37 ++ src/train/criterion/loss/__init__.py | 4 + src/train/criterion/loss/asymmetric_loss.py | 97 +++++ .../criterion/loss/binary_cross_entropy.py | 47 ++ src/train/criterion/loss/cross_entropy.py | 36 ++ src/train/criterion/loss/distillation_loss.py | 65 +++ src/train/criterion/loss/jsd.py | 39 ++ src/train/criterion/metrics/__init__.py | 0 src/train/criterion/metrics/accuracy.py | 19 + src/train/lr_scheduler_v2/__init__.py | 0 src/train/lr_scheduler_v2/knee_scheduler.py | 36 ++ src/train/lr_scheduler_v2/lr_scheduler.py | 46 ++ .../reduce_lr_on_plateau_lr_scheduler.py | 69 +++ .../transformer_lr_scheduler.py | 93 ++++ .../lr_scheduler_v2/tri_stage_lr_scheduler.py | 114 +++++ .../lr_scheduler_v2/warmup_lr_scheduler.py | 62 +++ .../warmup_reduce_lr_on_plateau_scheduler.py | 88 ++++ src/train/opt/__init__.py | 1 + src/train/opt/init_opt.py | 71 ++++ src/train/opt/opt_modules/__init__.py | 6 + src/train/opt/opt_modules/adafactor.py | 240 +++++++++++ src/train/opt/opt_modules/base_optimiser.py | 388 +++++++++++++++++ src/train/opt/opt_modules/ema.py | 14 + src/train/opt/opt_modules/lamb_opt.py | 153 +++++++ src/train/opt/opt_modules/lion_opt.py | 88 ++++ src/train/opt/opt_modules/lion_triton_opt.py | 97 +++++ src/train/opt/opt_modules/lookahead_opt.py | 153 +++++++ .../opt/opt_modules/lookahead_opt_old.py | 130 ++++++ src/train/opt/opt_modules/radam_opt.py | 119 ++++++ src/train/opt/opt_modules/weight_decay.py | 65 +++ src/train/train_naruto.py | 345 +++++++++++++++ src/utils/utils_file_dir.py | 151 +++++++ src/utils/utils_img.py | 172 ++++++++ src/utils/utils_log.py | 84 ++++ src/utils/utils_nn.py | 10 + src/utils/utils_train.py | 40 ++ src/utils/utils_types.py | 32 ++ 63 files changed, 5781 insertions(+) create mode 100644 .gitignore create mode 100644 env.yaml create mode 100644 help create mode 100644 in/config_files/file_dir.gin create mode 100644 in/config_files/input.gin create mode 100644 in/config_files/nn.gin create mode 100644 in/config_files/train.gin create mode 100644 src/__init__.py create mode 100644 src/conf/__init__.py create mode 100644 src/conf/base_conf.py create mode 100644 src/conf/nn_conf.py create mode 100644 src/conf/train_conf.py create mode 100644 src/ds/__init__.py create mode 100644 src/ds/gtauav_dl.py create mode 100644 src/ds/naruto_dl.py create mode 100644 src/models/__init__.py create mode 100644 src/models/bb/__init__.py create mode 100644 src/models/bb/coca/__init__.py create mode 100644 src/models/bb/coca/coca_model.py create mode 100644 src/models/bb/edgenext/edgenext_blocks.py create mode 100644 src/models/bb/edgenext/edgenext_model.py create mode 100644 src/models/bb/inception_next/__init__.py create mode 100644 src/models/bb/inception_next/inception_next_model.py create mode 100644 src/models/bb/load_models.py create mode 100644 src/models/bb/make_fe.py create mode 100644 src/train/criterion/__init__.py create mode 100644 src/train/criterion/init_loss.py create mode 100644 src/train/criterion/loss/__init__.py create mode 100644 src/train/criterion/loss/asymmetric_loss.py create mode 100644 src/train/criterion/loss/binary_cross_entropy.py create mode 100644 src/train/criterion/loss/cross_entropy.py create mode 100644 src/train/criterion/loss/distillation_loss.py create mode 100644 src/train/criterion/loss/jsd.py create mode 100644 src/train/criterion/metrics/__init__.py create mode 100644 src/train/criterion/metrics/accuracy.py create mode 100644 src/train/lr_scheduler_v2/__init__.py create mode 100644 src/train/lr_scheduler_v2/knee_scheduler.py create mode 100644 src/train/lr_scheduler_v2/lr_scheduler.py create mode 100644 src/train/lr_scheduler_v2/reduce_lr_on_plateau_lr_scheduler.py create mode 100644 src/train/lr_scheduler_v2/transformer_lr_scheduler.py create mode 100644 src/train/lr_scheduler_v2/tri_stage_lr_scheduler.py create mode 100644 src/train/lr_scheduler_v2/warmup_lr_scheduler.py create mode 100644 src/train/lr_scheduler_v2/warmup_reduce_lr_on_plateau_scheduler.py create mode 100644 src/train/opt/__init__.py create mode 100644 src/train/opt/init_opt.py create mode 100644 src/train/opt/opt_modules/__init__.py create mode 100644 src/train/opt/opt_modules/adafactor.py create mode 100644 src/train/opt/opt_modules/base_optimiser.py create mode 100644 src/train/opt/opt_modules/ema.py create mode 100644 src/train/opt/opt_modules/lamb_opt.py create mode 100644 src/train/opt/opt_modules/lion_opt.py create mode 100644 src/train/opt/opt_modules/lion_triton_opt.py create mode 100644 src/train/opt/opt_modules/lookahead_opt.py create mode 100644 src/train/opt/opt_modules/lookahead_opt_old.py create mode 100644 src/train/opt/opt_modules/radam_opt.py create mode 100644 src/train/opt/opt_modules/weight_decay.py create mode 100644 src/train/train_naruto.py create mode 100644 src/utils/utils_file_dir.py create mode 100644 src/utils/utils_img.py create mode 100644 src/utils/utils_log.py create mode 100644 src/utils/utils_nn.py create mode 100644 src/utils/utils_train.py create mode 100644 src/utils/utils_types.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f9d12ac --- /dev/null +++ b/.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 diff --git a/env.yaml b/env.yaml new file mode 100644 index 0000000..7212626 --- /dev/null +++ b/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 diff --git a/help b/help new file mode 100644 index 0000000..b4faea9 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/in/config_files/file_dir.gin b/in/config_files/file_dir.gin new file mode 100644 index 0000000..1be1ced --- /dev/null +++ b/in/config_files/file_dir.gin @@ -0,0 +1,3 @@ +# ============================================================================== + +FilesDirsInfo.path2data = "/home/uzver/Документы/datasets/" # \ No newline at end of file diff --git a/in/config_files/input.gin b/in/config_files/input.gin new file mode 100644 index 0000000..0bd0729 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/in/config_files/nn.gin b/in/config_files/nn.gin new file mode 100644 index 0000000..e597cc9 --- /dev/null +++ b/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' \ No newline at end of file diff --git a/in/config_files/train.gin b/in/config_files/train.gin new file mode 100644 index 0000000..a53a01f --- /dev/null +++ b/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/' + diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/conf/__init__.py b/src/conf/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/conf/base_conf.py b/src/conf/base_conf.py new file mode 100644 index 0000000..b5cca40 --- /dev/null +++ b/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 diff --git a/src/conf/nn_conf.py b/src/conf/nn_conf.py new file mode 100644 index 0000000..19acaea --- /dev/null +++ b/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 diff --git a/src/conf/train_conf.py b/src/conf/train_conf.py new file mode 100644 index 0000000..1ff54af --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/ds/__init__.py b/src/ds/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/ds/gtauav_dl.py b/src/ds/gtauav_dl.py new file mode 100644 index 0000000..abedbef --- /dev/null +++ b/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() \ No newline at end of file diff --git a/src/ds/naruto_dl.py b/src/ds/naruto_dl.py new file mode 100644 index 0000000..7b18708 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/src/models/__init__.py b/src/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/models/bb/__init__.py b/src/models/bb/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/models/bb/coca/__init__.py b/src/models/bb/coca/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/models/bb/coca/coca_model.py b/src/models/bb/coca/coca_model.py new file mode 100644 index 0000000..ba8bd0e --- /dev/null +++ b/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() \ No newline at end of file diff --git a/src/models/bb/edgenext/edgenext_blocks.py b/src/models/bb/edgenext/edgenext_blocks.py new file mode 100644 index 0000000..98961bc --- /dev/null +++ b/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'} \ No newline at end of file diff --git a/src/models/bb/edgenext/edgenext_model.py b/src/models/bb/edgenext/edgenext_model.py new file mode 100644 index 0000000..57e7abe --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/models/bb/inception_next/__init__.py b/src/models/bb/inception_next/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/models/bb/inception_next/inception_next_model.py b/src/models/bb/inception_next/inception_next_model.py new file mode 100644 index 0000000..669c8da --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/models/bb/load_models.py b/src/models/bb/load_models.py new file mode 100644 index 0000000..ce9339c --- /dev/null +++ b/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() + + + + + diff --git a/src/models/bb/make_fe.py b/src/models/bb/make_fe.py new file mode 100644 index 0000000..7544995 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/src/train/criterion/__init__.py b/src/train/criterion/__init__.py new file mode 100644 index 0000000..4a71414 --- /dev/null +++ b/src/train/criterion/__init__.py @@ -0,0 +1 @@ +from src.train.criterion.init_loss import get_loss_func \ No newline at end of file diff --git a/src/train/criterion/init_loss.py b/src/train/criterion/init_loss.py new file mode 100644 index 0000000..df70858 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/criterion/loss/__init__.py b/src/train/criterion/loss/__init__.py new file mode 100644 index 0000000..ea7f15f --- /dev/null +++ b/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 diff --git a/src/train/criterion/loss/asymmetric_loss.py b/src/train/criterion/loss/asymmetric_loss.py new file mode 100644 index 0000000..96a9778 --- /dev/null +++ b/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 diff --git a/src/train/criterion/loss/binary_cross_entropy.py b/src/train/criterion/loss/binary_cross_entropy.py new file mode 100644 index 0000000..ed76c1e --- /dev/null +++ b/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) diff --git a/src/train/criterion/loss/cross_entropy.py b/src/train/criterion/loss/cross_entropy.py new file mode 100644 index 0000000..8519810 --- /dev/null +++ b/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() diff --git a/src/train/criterion/loss/distillation_loss.py b/src/train/criterion/loss/distillation_loss.py new file mode 100644 index 0000000..0fc828d --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/criterion/loss/jsd.py b/src/train/criterion/loss/jsd.py new file mode 100644 index 0000000..dd64e15 --- /dev/null +++ b/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 diff --git a/src/train/criterion/metrics/__init__.py b/src/train/criterion/metrics/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/train/criterion/metrics/accuracy.py b/src/train/criterion/metrics/accuracy.py new file mode 100644 index 0000000..b2dc1c3 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/lr_scheduler_v2/__init__.py b/src/train/lr_scheduler_v2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/train/lr_scheduler_v2/knee_scheduler.py b/src/train/lr_scheduler_v2/knee_scheduler.py new file mode 100644 index 0000000..29f1698 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/src/train/lr_scheduler_v2/lr_scheduler.py b/src/train/lr_scheduler_v2/lr_scheduler.py new file mode 100644 index 0000000..76c7a29 --- /dev/null +++ b/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'] diff --git a/src/train/lr_scheduler_v2/reduce_lr_on_plateau_lr_scheduler.py b/src/train/lr_scheduler_v2/reduce_lr_on_plateau_lr_scheduler.py new file mode 100644 index 0000000..b4e96ac --- /dev/null +++ b/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 diff --git a/src/train/lr_scheduler_v2/transformer_lr_scheduler.py b/src/train/lr_scheduler_v2/transformer_lr_scheduler.py new file mode 100644 index 0000000..d65c70a --- /dev/null +++ b/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 diff --git a/src/train/lr_scheduler_v2/tri_stage_lr_scheduler.py b/src/train/lr_scheduler_v2/tri_stage_lr_scheduler.py new file mode 100644 index 0000000..f753e4c --- /dev/null +++ b/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 diff --git a/src/train/lr_scheduler_v2/warmup_lr_scheduler.py b/src/train/lr_scheduler_v2/warmup_lr_scheduler.py new file mode 100644 index 0000000..3117705 --- /dev/null +++ b/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 diff --git a/src/train/lr_scheduler_v2/warmup_reduce_lr_on_plateau_scheduler.py b/src/train/lr_scheduler_v2/warmup_reduce_lr_on_plateau_scheduler.py new file mode 100644 index 0000000..f8fb303 --- /dev/null +++ b/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() diff --git a/src/train/opt/__init__.py b/src/train/opt/__init__.py new file mode 100644 index 0000000..5454a52 --- /dev/null +++ b/src/train/opt/__init__.py @@ -0,0 +1 @@ +from src.train.opt.init_opt import create_optimizer \ No newline at end of file diff --git a/src/train/opt/init_opt.py b/src/train/opt/init_opt.py new file mode 100644 index 0000000..99e4cf0 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/opt/opt_modules/__init__.py b/src/train/opt/opt_modules/__init__.py new file mode 100644 index 0000000..6a5e87e --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/opt/opt_modules/adafactor.py b/src/train/opt/opt_modules/adafactor.py new file mode 100644 index 0000000..346482b --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/opt/opt_modules/base_optimiser.py b/src/train/opt/opt_modules/base_optimiser.py new file mode 100644 index 0000000..8d1c15f --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/opt/opt_modules/ema.py b/src/train/opt/opt_modules/ema.py new file mode 100644 index 0000000..2858b66 --- /dev/null +++ b/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 `_ + 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) \ No newline at end of file diff --git a/src/train/opt/opt_modules/lamb_opt.py b/src/train/opt/opt_modules/lamb_opt.py new file mode 100644 index 0000000..7066645 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/opt/opt_modules/lion_opt.py b/src/train/opt/opt_modules/lion_opt.py new file mode 100644 index 0000000..f6b5de7 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/opt/opt_modules/lion_triton_opt.py b/src/train/opt/opt_modules/lion_triton_opt.py new file mode 100644 index 0000000..4c7d785 --- /dev/null +++ b/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 + ) \ No newline at end of file diff --git a/src/train/opt/opt_modules/lookahead_opt.py b/src/train/opt/opt_modules/lookahead_opt.py new file mode 100644 index 0000000..bf4649c --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/opt/opt_modules/lookahead_opt_old.py b/src/train/opt/opt_modules/lookahead_opt_old.py new file mode 100644 index 0000000..cf80f67 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/opt/opt_modules/radam_opt.py b/src/train/opt/opt_modules/radam_opt.py new file mode 100644 index 0000000..e71f967 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/opt/opt_modules/weight_decay.py b/src/train/opt/opt_modules/weight_decay.py new file mode 100644 index 0000000..0ab1d90 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/train/train_naruto.py b/src/train/train_naruto.py new file mode 100644 index 0000000..c8c70a2 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/src/utils/utils_file_dir.py b/src/utils/utils_file_dir.py new file mode 100644 index 0000000..ece996a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/utils/utils_img.py b/src/utils/utils_img.py new file mode 100644 index 0000000..6d93449 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/utils/utils_log.py b/src/utils/utils_log.py new file mode 100644 index 0000000..37d1504 --- /dev/null +++ b/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']) diff --git a/src/utils/utils_nn.py b/src/utils/utils_nn.py new file mode 100644 index 0000000..74a9939 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/utils/utils_train.py b/src/utils/utils_train.py new file mode 100644 index 0000000..fa71bc7 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/utils/utils_types.py b/src/utils/utils_types.py new file mode 100644 index 0000000..f1295c1 --- /dev/null +++ b/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