Integration of I-ViT operators
Context
Implementation of the 3 I-ViT operators: ShiftMax, ShiftGELU and ILayerNorm, which are quantifiable arithmetic approximations of SoftMax, GELU and LayerNorm. The forward and backward operators are implemented to enable quantization-aware-training of networks containing these layers (notably Transformers). Unit tests are also implemented to illustrate how these operators work, and to measure the difference between these approximators and the original operators.
Modified files
-
include/aidge/backend/cuda.hpp
: Import of ShiftMax, ShiftGELU and ILayerNorm;
Added files
-
src/operator/ShiftMaxImpl.cpp
andinclude/aidge/backend/cuda/operator/ShiftMaxImpl.hpp
: Definition of the forward and backward for the ShiftMax method (highly inspired by ReLU implementation); -
src/operator/ShiftMaxImpl_CUDA_kernels.cu
andinclude/aidge/backend/cuda/operator/ShiftMaxImpl_CUDA_kernels.hpp
: Implementation of the CUDA kernels which compute the ShiftMax operation. This CUDA code is conform to the Python implementation of I-ViT; -
unit_tests/Test_ShiftMaxImpl.cpp
: Unit test that illustrate how the forward and backward function of ShiftMax work and compare it with SoftMax; -
src/operator/ShiftGELUImpl.cpp
andinclude/aidge/backend/cuda/operator/ShiftGELUImpl.hpp
: Definition of the forward and backward for the ShiftGELU method (highly inspired by ReLU implementation); -
src/operator/ShiftGELUImpl_CUDA_kernels.cu
andinclude/aidge/backend/cuda/operator/ShiftGELUImpl_CUDA_kernels.hpp
: Implementation of the CUDA kernels which compute the ShiftGELU operation. This CUDA code is conform to the Python implementation of I-ViT; -
unit_tests/Test_ShiftGELUImpl.cpp
: Unit test that illustrate how the forward and backward function of ShiftGELU work and compare it with SoftMax; -
src/operator/ILayerNormImpl.cpp
andinclude/aidge/backend/cuda/operator/ILayerNormImpl.hpp
: Definition of the forward and backward for the ILayerNorm method (highly inspired by FC implementation); -
src/operator/ILayerNormImpl_CUDA_kernels.cu
andinclude/aidge/backend/cuda/operator/ILayerNormImpl_CUDA_kernels.hpp
: Implementation of the CUDA kernels which compute the ILayerNorm operation. This CUDA code is conform to the Python implementation of I-ViT; -
unit_tests/Test_IlayerNormImpl.cpp
: Unit test that illustrate how the forward and backward function of ILayerNorm work;
TODO
Merge into main if pipeline successful
-
NOT DONE -
DONE -
TO DO