├── nn ├── backends │ ├── __init__.py │ ├── __pycache__ │ │ ├── thnn.cpython-36.pyc │ │ └── __init__.cpython-36.pyc │ └── thnn.py ├── .DS_Store ├── __pycache__ │ ├── _VF.cpython-36.pyc │ ├── cpp.cpython-36.pyc │ ├── grad.cpython-36.pyc │ ├── init.cpython-36.pyc │ ├── __init__.cpython-36.pyc │ ├── _reduction.cpython-36.pyc │ ├── functional.cpython-36.pyc │ └── parameter.cpython-36.pyc ├── qat │ ├── __init__.py │ ├── __pycache__ │ │ └── __init__.cpython-36.pyc │ └── modules │ │ ├── __pycache__ │ │ ├── conv.cpython-36.pyc │ │ ├── linear.cpython-36.pyc │ │ └── __init__.cpython-36.pyc │ │ ├── __init__.py │ │ ├── linear.py │ │ └── conv.py ├── modules │ ├── __pycache__ │ │ ├── rnn.cpython-36.pyc │ │ ├── conv.cpython-36.pyc │ │ ├── fold.cpython-36.pyc │ │ ├── lala.cpython-36.pyc │ │ ├── loss.cpython-36.pyc │ │ ├── utils.cpython-36.pyc │ │ ├── __init__.cpython-36.pyc │ │ ├── adaptive.cpython-36.pyc │ │ ├── distance.cpython-36.pyc │ │ ├── dropout.cpython-36.pyc │ │ ├── flatten.cpython-36.pyc │ │ ├── linear.cpython-36.pyc │ │ ├── module.cpython-36.pyc │ │ ├── padding.cpython-36.pyc │ │ ├── pooling.cpython-36.pyc │ │ ├── sparse.cpython-36.pyc │ │ ├── _functions.cpython-36.pyc │ │ ├── activation.cpython-36.pyc │ │ ├── batchnorm.cpython-36.pyc │ │ ├── container.cpython-36.pyc │ │ ├── upsampling.cpython-36.pyc │ │ ├── instancenorm.cpython-36.pyc │ │ ├── normalization.cpython-36.pyc │ │ ├── pixelshuffle.cpython-36.pyc │ │ └── transformer.cpython-36.pyc │ ├── pixelshuffle.pyi │ ├── flatten.pyi │ ├── instancenorm.pyi │ ├── utils.py │ ├── distance.pyi │ ├── flatten.py │ ├── upsampling.pyi │ ├── adaptive.pyi │ ├── fold.pyi │ ├── dropout.pyi │ ├── linear.pyi │ ├── batchnorm.pyi │ ├── pixelshuffle.py │ ├── normalization.pyi │ ├── padding.pyi │ ├── sparse.pyi │ ├── distance.py │ ├── __init__.pyi │ ├── container.pyi │ ├── __init__.py │ ├── module.pyi │ ├── conv.pyi │ ├── rnn.pyi │ ├── activation.pyi │ ├── linear.py │ ├── loss.pyi │ └── dropout.py ├── utils │ ├── __pycache__ │ │ ├── prune.cpython-36.pyc │ │ ├── rnn.cpython-36.pyc │ │ ├── fusion.cpython-36.pyc │ │ ├── __init__.cpython-36.pyc │ │ ├── clip_grad.cpython-36.pyc │ │ ├── weight_norm.cpython-36.pyc │ │ ├── spectral_norm.cpython-36.pyc │ │ └── convert_parameters.cpython-36.pyc │ ├── __init__.py │ ├── fusion.py │ ├── clip_grad.py │ ├── convert_parameters.py │ └── weight_norm.py ├── intrinsic │ ├── __pycache__ │ │ └── __init__.cpython-36.pyc │ ├── modules │ │ ├── __pycache__ │ │ │ ├── fused.cpython-36.pyc │ │ │ └── __init__.cpython-36.pyc │ │ ├── __init__.py │ │ └── fused.py │ ├── qat │ │ ├── __pycache__ │ │ │ └── __init__.cpython-36.pyc │ │ ├── modules │ │ │ ├── __pycache__ │ │ │ │ ├── __init__.cpython-36.pyc │ │ │ │ ├── conv_fused.cpython-36.pyc │ │ │ │ └── linear_relu.cpython-36.pyc │ │ │ ├── __init__.py │ │ │ └── linear_relu.py │ │ └── __init__.py │ ├── quantized │ │ ├── __pycache__ │ │ │ └── __init__.cpython-36.pyc │ │ ├── modules │ │ │ ├── __pycache__ │ │ │ │ ├── __init__.cpython-36.pyc │ │ │ │ ├── conv_relu.cpython-36.pyc │ │ │ │ └── linear_relu.cpython-36.pyc │ │ │ ├── __init__.py │ │ │ ├── linear_relu.py │ │ │ └── conv_relu.py │ │ └── __init__.py │ └── __init__.py ├── parallel │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── replicate.cpython-36.pyc │ │ ├── _functions.cpython-36.pyc │ │ ├── distributed.cpython-36.pyc │ │ ├── data_parallel.cpython-36.pyc │ │ ├── parallel_apply.cpython-36.pyc │ │ └── scatter_gather.cpython-36.pyc │ ├── common_types.pyi │ ├── replicate.pyi │ ├── parallel_apply.pyi │ ├── __init__.pyi │ ├── __init__.py │ ├── data_parallel.pyi │ ├── scatter_gather.pyi │ ├── distributed.pyi │ ├── scatter_gather.py │ ├── parallel_apply.py │ ├── _functions.py │ └── replicate.py ├── quantized │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ └── functional.cpython-36.pyc │ ├── dynamic │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ └── __init__.cpython-36.pyc │ │ └── modules │ │ │ ├── __pycache__ │ │ │ ├── rnn.cpython-36.pyc │ │ │ ├── linear.cpython-36.pyc │ │ │ └── __init__.cpython-36.pyc │ │ │ ├── __init__.py │ │ │ └── linear.py │ ├── modules │ │ ├── __pycache__ │ │ │ ├── conv.cpython-36.pyc │ │ │ ├── utils.cpython-36.pyc │ │ │ ├── __init__.cpython-36.pyc │ │ │ ├── linear.cpython-36.pyc │ │ │ ├── activation.cpython-36.pyc │ │ │ └── functional_modules.cpython-36.pyc │ │ ├── utils.py │ │ ├── activation.py │ │ ├── __init__.py │ │ └── functional_modules.py │ └── __init__.py ├── __init__.py ├── parameter.pyi ├── __init__.pyi ├── _VF.py ├── common_types.pyi ├── _reduction.py ├── parameter.py └── cpp.py ├── quantization ├── __pycache__ │ ├── stubs.cpython-36.pyc │ ├── __init__.cpython-36.pyc │ ├── observer.cpython-36.pyc │ ├── qconfig.cpython-36.pyc │ ├── quantize.cpython-36.pyc │ ├── fake_quantize.cpython-36.pyc │ ├── fuse_modules.cpython-36.pyc │ ├── _quantize_script.cpython-36.pyc │ └── default_mappings.cpython-36.pyc ├── __init__.py ├── stubs.py ├── default_mappings.py ├── _quantize_script.py ├── qconfig.py ├── fuse_modules.py └── fake_quantize.py ├── scheduler ├── Makefile ├── lalarand.h ├── lalarand_fn.h └── lalarand.c └── README.md /nn/backends/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /nn/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/.DS_Store -------------------------------------------------------------------------------- /nn/__pycache__/_VF.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/__pycache__/_VF.cpython-36.pyc -------------------------------------------------------------------------------- /nn/__pycache__/cpp.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/__pycache__/cpp.cpython-36.pyc -------------------------------------------------------------------------------- /nn/__pycache__/grad.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/__pycache__/grad.cpython-36.pyc -------------------------------------------------------------------------------- /nn/__pycache__/init.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/__pycache__/init.cpython-36.pyc -------------------------------------------------------------------------------- /nn/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/qat/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | from .modules import * 3 | -------------------------------------------------------------------------------- /nn/__pycache__/_reduction.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/__pycache__/_reduction.cpython-36.pyc -------------------------------------------------------------------------------- /nn/__pycache__/functional.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/__pycache__/functional.cpython-36.pyc -------------------------------------------------------------------------------- /nn/__pycache__/parameter.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/__pycache__/parameter.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/rnn.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/rnn.cpython-36.pyc -------------------------------------------------------------------------------- /nn/utils/__pycache__/prune.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/utils/__pycache__/prune.cpython-36.pyc -------------------------------------------------------------------------------- /nn/utils/__pycache__/rnn.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/utils/__pycache__/rnn.cpython-36.pyc -------------------------------------------------------------------------------- /nn/backends/__pycache__/thnn.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/backends/__pycache__/thnn.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/conv.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/conv.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/fold.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/fold.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/lala.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/lala.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/loss.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /nn/qat/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/qat/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/utils/__pycache__/fusion.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/utils/__pycache__/fusion.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/adaptive.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/adaptive.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/distance.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/distance.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/dropout.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/dropout.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/flatten.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/flatten.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/linear.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/linear.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/module.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/module.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/padding.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/padding.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/pooling.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/pooling.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/sparse.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/sparse.cpython-36.pyc -------------------------------------------------------------------------------- /nn/qat/modules/__pycache__/conv.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/qat/modules/__pycache__/conv.cpython-36.pyc -------------------------------------------------------------------------------- /nn/utils/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/utils/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/utils/__pycache__/clip_grad.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/utils/__pycache__/clip_grad.cpython-36.pyc -------------------------------------------------------------------------------- /quantization/__pycache__/stubs.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/quantization/__pycache__/stubs.cpython-36.pyc -------------------------------------------------------------------------------- /nn/backends/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/backends/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/backends/thnn.py: -------------------------------------------------------------------------------- 1 | # this is for historical pickle deserilaization, it is not used otherwise 2 | 3 | def _get_thnn_function_backend(): 4 | pass 5 | -------------------------------------------------------------------------------- /nn/intrinsic/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/_functions.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/_functions.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/activation.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/activation.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/batchnorm.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/batchnorm.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/container.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/container.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/upsampling.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/upsampling.cpython-36.pyc -------------------------------------------------------------------------------- /nn/parallel/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/parallel/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/parallel/__pycache__/replicate.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/parallel/__pycache__/replicate.cpython-36.pyc -------------------------------------------------------------------------------- /nn/qat/modules/__pycache__/linear.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/qat/modules/__pycache__/linear.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/utils/__pycache__/weight_norm.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/utils/__pycache__/weight_norm.cpython-36.pyc -------------------------------------------------------------------------------- /quantization/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/quantization/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /quantization/__pycache__/observer.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/quantization/__pycache__/observer.cpython-36.pyc -------------------------------------------------------------------------------- /quantization/__pycache__/qconfig.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/quantization/__pycache__/qconfig.cpython-36.pyc -------------------------------------------------------------------------------- /quantization/__pycache__/quantize.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/quantization/__pycache__/quantize.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/instancenorm.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/instancenorm.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/normalization.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/normalization.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/pixelshuffle.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/pixelshuffle.cpython-36.pyc -------------------------------------------------------------------------------- /nn/modules/__pycache__/transformer.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/modules/__pycache__/transformer.cpython-36.pyc -------------------------------------------------------------------------------- /nn/parallel/__pycache__/_functions.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/parallel/__pycache__/_functions.cpython-36.pyc -------------------------------------------------------------------------------- /nn/parallel/__pycache__/distributed.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/parallel/__pycache__/distributed.cpython-36.pyc -------------------------------------------------------------------------------- /nn/qat/modules/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/qat/modules/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/__pycache__/functional.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/__pycache__/functional.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/dynamic/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from __future__ import absolute_import, division, print_function, unicode_literals 3 | 4 | from .modules import * 5 | -------------------------------------------------------------------------------- /nn/utils/__pycache__/spectral_norm.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/utils/__pycache__/spectral_norm.cpython-36.pyc -------------------------------------------------------------------------------- /nn/intrinsic/modules/__pycache__/fused.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/modules/__pycache__/fused.cpython-36.pyc -------------------------------------------------------------------------------- /nn/intrinsic/qat/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/qat/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/parallel/__pycache__/data_parallel.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/parallel/__pycache__/data_parallel.cpython-36.pyc -------------------------------------------------------------------------------- /nn/parallel/__pycache__/parallel_apply.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/parallel/__pycache__/parallel_apply.cpython-36.pyc -------------------------------------------------------------------------------- /nn/parallel/__pycache__/scatter_gather.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/parallel/__pycache__/scatter_gather.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/modules/__pycache__/conv.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/modules/__pycache__/conv.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/modules/__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/modules/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /quantization/__pycache__/fake_quantize.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/quantization/__pycache__/fake_quantize.cpython-36.pyc -------------------------------------------------------------------------------- /quantization/__pycache__/fuse_modules.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/quantization/__pycache__/fuse_modules.cpython-36.pyc -------------------------------------------------------------------------------- /nn/__init__.py: -------------------------------------------------------------------------------- 1 | from .modules import * 2 | from .parameter import Parameter 3 | from .parallel import DataParallel 4 | from . import init 5 | from . import utils 6 | -------------------------------------------------------------------------------- /nn/intrinsic/modules/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/modules/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/dynamic/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/dynamic/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/modules/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/modules/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/modules/__pycache__/linear.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/modules/__pycache__/linear.cpython-36.pyc -------------------------------------------------------------------------------- /nn/utils/__pycache__/convert_parameters.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/utils/__pycache__/convert_parameters.cpython-36.pyc -------------------------------------------------------------------------------- /quantization/__pycache__/_quantize_script.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/quantization/__pycache__/_quantize_script.cpython-36.pyc -------------------------------------------------------------------------------- /quantization/__pycache__/default_mappings.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/quantization/__pycache__/default_mappings.cpython-36.pyc -------------------------------------------------------------------------------- /nn/intrinsic/quantized/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/quantized/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/parallel/common_types.pyi: -------------------------------------------------------------------------------- 1 | from typing import Union, Sequence 2 | from ... import device 3 | 4 | _device_t = Union[int, device] 5 | _devices_t = Sequence[_device_t] 6 | -------------------------------------------------------------------------------- /nn/quantized/modules/__pycache__/activation.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/modules/__pycache__/activation.cpython-36.pyc -------------------------------------------------------------------------------- /nn/intrinsic/qat/modules/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/qat/modules/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/dynamic/modules/__pycache__/rnn.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/dynamic/modules/__pycache__/rnn.cpython-36.pyc -------------------------------------------------------------------------------- /nn/intrinsic/qat/modules/__pycache__/conv_fused.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/qat/modules/__pycache__/conv_fused.cpython-36.pyc -------------------------------------------------------------------------------- /nn/intrinsic/qat/modules/__pycache__/linear_relu.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/qat/modules/__pycache__/linear_relu.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/dynamic/modules/__pycache__/linear.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/dynamic/modules/__pycache__/linear.cpython-36.pyc -------------------------------------------------------------------------------- /nn/intrinsic/quantized/modules/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/quantized/modules/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/dynamic/modules/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/dynamic/modules/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /nn/quantized/modules/__pycache__/functional_modules.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/quantized/modules/__pycache__/functional_modules.cpython-36.pyc -------------------------------------------------------------------------------- /nn/intrinsic/quantized/modules/__pycache__/conv_relu.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/quantized/modules/__pycache__/conv_relu.cpython-36.pyc -------------------------------------------------------------------------------- /nn/parameter.pyi: -------------------------------------------------------------------------------- 1 | from .. import Tensor 2 | import builtins 3 | 4 | class Parameter(Tensor): 5 | def __init__(self, data: Tensor, requires_grad: builtins.bool): ... 6 | 7 | ... 8 | -------------------------------------------------------------------------------- /nn/intrinsic/quantized/modules/__pycache__/linear_relu.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fredrickang/LaLaRAND/HEAD/nn/intrinsic/quantized/modules/__pycache__/linear_relu.cpython-36.pyc -------------------------------------------------------------------------------- /nn/qat/modules/__init__.py: -------------------------------------------------------------------------------- 1 | # @lint-ignore-every PYTHON3COMPATIMPORTS 2 | 3 | from .linear import Linear 4 | from .conv import Conv2d 5 | 6 | __all__ = [ 7 | 'Linear', 8 | 'Conv2d', 9 | ] 10 | -------------------------------------------------------------------------------- /nn/quantized/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | from __future__ import unicode_literals 5 | 6 | from .modules import * 7 | -------------------------------------------------------------------------------- /nn/quantized/dynamic/modules/__init__.py: -------------------------------------------------------------------------------- 1 | # @lint-ignore-every PYTHON3COMPATIMPORTS 2 | 3 | from .linear import Linear 4 | from .rnn import LSTM 5 | 6 | __all__ = [ 7 | 'Linear', 8 | 'LSTM', 9 | ] 10 | -------------------------------------------------------------------------------- /nn/__init__.pyi: -------------------------------------------------------------------------------- 1 | from .modules import * 2 | from .parameter import Parameter as Parameter 3 | from .parallel import DataParallel as DataParallel 4 | from . import init as init 5 | from . import utils as utils 6 | from . import functional as functional 7 | -------------------------------------------------------------------------------- /nn/intrinsic/quantized/modules/__init__.py: -------------------------------------------------------------------------------- 1 | # @lint-ignore-every PYTHON3COMPATIMPORTS 2 | 3 | from .linear_relu import LinearReLU 4 | from .conv_relu import ConvReLU2d, ConvReLU3d 5 | 6 | __all__ = [ 7 | 'LinearReLU', 8 | 'ConvReLU2d', 9 | 'ConvReLU3d', 10 | ] 11 | -------------------------------------------------------------------------------- /nn/intrinsic/quantized/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | from .modules import LinearReLU 3 | from .modules import ConvReLU2d, ConvReLU3d 4 | 5 | __all__ = [ 6 | 'LinearReLU', 7 | 'ConvReLU2d', 8 | 'ConvReLU3d', 9 | ] 10 | -------------------------------------------------------------------------------- /nn/parallel/replicate.pyi: -------------------------------------------------------------------------------- 1 | from typing import List, Union, Sequence, TypeVar 2 | from ..modules import Module 3 | from .common_types import _devices_t 4 | 5 | T = TypeVar('T') 6 | 7 | 8 | def replicate(network: Module[T], devices: Union[_devices_t, Sequence[_devices_t]], detach: bool = ...) -> List[ 9 | Module[T]]: ... 10 | -------------------------------------------------------------------------------- /scheduler/Makefile: -------------------------------------------------------------------------------- 1 | CC=g++ 2 | CFLAGS=-g -Wall 3 | OBJS=lalarand.o lalarand_fn.o 4 | TARGET=lalarand 5 | 6 | $(TARGET) : $(OBJS) 7 | $(CC) -o $@ $(OBJS) 8 | 9 | lalarand.o: lalarand.h lalarand_fn.h lalarand.c 10 | lalarand_fn.o: lalarand_fn.h lalarand.h lalarand_fn.c 11 | 12 | clean: 13 | rm -f *.o 14 | rm -f $(TARGET) 15 | -------------------------------------------------------------------------------- /nn/parallel/parallel_apply.pyi: -------------------------------------------------------------------------------- 1 | from typing import Any, Optional, Sequence, List 2 | from .common_types import _devices_t 3 | from ..modules import Module 4 | 5 | 6 | def parallel_apply(modules: Sequence[Module], inputs: Sequence[Any], kwargs_tup: Optional[Any] = ..., 7 | devices: Optional[_devices_t] = ...) -> List[Any]: ... 8 | -------------------------------------------------------------------------------- /nn/parallel/__init__.pyi: -------------------------------------------------------------------------------- 1 | from .data_parallel import DataParallel as DataParallel, data_parallel as data_parallel 2 | from .distributed import DistributedDataParallel as DistributedDataParallel 3 | from .parallel_apply import parallel_apply as parallel_apply 4 | from .replicate import replicate as replicate 5 | from .scatter_gather import gather as gather, scatter as scatter 6 | -------------------------------------------------------------------------------- /nn/_VF.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import sys 3 | import types 4 | 5 | 6 | class VFModule(types.ModuleType): 7 | def __init__(self, name): 8 | super(VFModule, self).__init__(name) 9 | self.vf = torch._C._VariableFunctions 10 | 11 | def __getattr__(self, attr): 12 | return getattr(self.vf, attr) 13 | 14 | sys.modules[__name__] = VFModule(__name__) 15 | -------------------------------------------------------------------------------- /nn/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from . import rnn 2 | from .clip_grad import clip_grad_norm, clip_grad_norm_, clip_grad_value_ 3 | from .weight_norm import weight_norm, remove_weight_norm 4 | from .convert_parameters import parameters_to_vector, vector_to_parameters 5 | from .spectral_norm import spectral_norm, remove_spectral_norm 6 | from .fusion import fuse_conv_bn_eval, fuse_conv_bn_weights 7 | -------------------------------------------------------------------------------- /nn/modules/pixelshuffle.pyi: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from ... import Tensor 3 | 4 | 5 | class PixelShuffle(Module): 6 | upscale_factor: int = ... 7 | 8 | def __init__(self, upscale_factor: int) -> None: ... 9 | 10 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 11 | 12 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 13 | -------------------------------------------------------------------------------- /nn/intrinsic/__init__.py: -------------------------------------------------------------------------------- 1 | # @lint-ignore-every PYTHON3COMPATIMPORTS 2 | 3 | from .modules import ConvBn2d 4 | from .modules import ConvBnReLU2d 5 | from .modules import ConvReLU2d 6 | from .modules import ConvReLU3d 7 | from .modules import LinearReLU 8 | 9 | __all__ = [ 10 | 'ConvBn2d', 11 | 'ConvBnReLU2d', 12 | 'ConvReLU2d', 13 | 'ConvReLU3d', 14 | 'LinearReLU', 15 | ] 16 | -------------------------------------------------------------------------------- /nn/intrinsic/modules/__init__.py: -------------------------------------------------------------------------------- 1 | # @lint-ignore-every PYTHON3COMPATIMPORTS 2 | 3 | from .fused import ConvBn2d 4 | from .fused import ConvBnReLU2d 5 | from .fused import ConvReLU2d 6 | from .fused import ConvReLU3d 7 | from .fused import LinearReLU 8 | 9 | __all__ = [ 10 | 'ConvBn2d', 11 | 'ConvBnReLU2d', 12 | 'ConvReLU2d', 13 | 'ConvReLU3d', 14 | 'LinearReLU', 15 | ] 16 | -------------------------------------------------------------------------------- /nn/modules/flatten.pyi: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | from .module import Module 3 | 4 | class Flatten(Module): 5 | __constants__: Any = ... 6 | start_dim: Any = ... 7 | end_dim: Any = ... 8 | def __init__(self, start_dim: int = ..., end_dim: int = ...) -> None: ... 9 | def forward(self, input: Any): ... # type: ignore 10 | def __call__(self, input: Any): ... # type: ignore 11 | -------------------------------------------------------------------------------- /nn/intrinsic/qat/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | 3 | from .linear_relu import LinearReLU 4 | from .conv_fused import ConvBn2d, ConvBnReLU2d, ConvReLU2d, update_bn_stats, freeze_bn_stats 5 | 6 | __all__ = [ 7 | 'LinearReLU', 8 | 'ConvReLU2d', 9 | 'ConvBn2d', 10 | 'ConvBnReLU2d', 11 | 'update_bn_stats', 12 | 'freeze_bn_stats' 13 | ] 14 | -------------------------------------------------------------------------------- /nn/intrinsic/qat/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | from .modules import LinearReLU 3 | from .modules import ConvReLU2d 4 | from .modules import ConvBn2d 5 | from .modules import ConvBnReLU2d 6 | from .modules import update_bn_stats, freeze_bn_stats 7 | 8 | __all__ = [ 9 | 'ConvBn2d', 10 | 'ConvBnReLU2d', 11 | 'ConvReLU2d', 12 | 'LinearReLU', 13 | 'update_bn_stats', 14 | 'freeze_bn_stats' 15 | ] 16 | -------------------------------------------------------------------------------- /nn/modules/instancenorm.pyi: -------------------------------------------------------------------------------- 1 | from ... import Tensor 2 | from .batchnorm import _BatchNorm 3 | 4 | 5 | class _InstanceNorm(_BatchNorm): 6 | def __init__(self, num_features: int, eps: float = ..., momentum: float = ..., affine: bool = ..., 7 | track_running_stats: bool = ...) -> None: ... 8 | 9 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 10 | 11 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 12 | 13 | 14 | class InstanceNorm1d(_InstanceNorm): ... 15 | 16 | 17 | class InstanceNorm2d(_InstanceNorm): ... 18 | 19 | 20 | class InstanceNorm3d(_InstanceNorm): ... 21 | -------------------------------------------------------------------------------- /nn/parallel/__init__.py: -------------------------------------------------------------------------------- 1 | from .parallel_apply import parallel_apply 2 | from .replicate import replicate 3 | from .data_parallel import DataParallel, data_parallel 4 | from .scatter_gather import scatter, gather 5 | from .distributed import DistributedDataParallel 6 | 7 | __all__ = ['replicate', 'scatter', 'parallel_apply', 'gather', 'data_parallel', 8 | 'DataParallel', 'DistributedDataParallel'] 9 | 10 | def DistributedDataParallelCPU(*args, **kwargs): 11 | import warnings 12 | warnings.warn("torch.nn.parallel.DistributedDataParallelCPU is deprecated, " 13 | "please use torch.nn.parallel.DistributedDataParallel instead.") 14 | return DistributedDataParallel(*args, **kwargs) 15 | -------------------------------------------------------------------------------- /nn/modules/utils.py: -------------------------------------------------------------------------------- 1 | from torch._six import container_abcs 2 | from itertools import repeat 3 | 4 | 5 | def _ntuple(n): 6 | def parse(x): 7 | if isinstance(x, container_abcs.Iterable): 8 | return x 9 | return tuple(repeat(x, n)) 10 | return parse 11 | 12 | _single = _ntuple(1) 13 | _pair = _ntuple(2) 14 | _triple = _ntuple(3) 15 | _quadruple = _ntuple(4) 16 | 17 | 18 | def _list_with_default(out_size, defaults): 19 | if isinstance(out_size, int): 20 | return out_size 21 | if len(defaults) <= len(out_size): 22 | raise ValueError('Input dimension should be at least {}'.format(len(out_size) + 1)) 23 | return [v if v is not None else d for v, d in zip(out_size, defaults[-len(out_size):])] 24 | -------------------------------------------------------------------------------- /nn/modules/distance.pyi: -------------------------------------------------------------------------------- 1 | from ... import Tensor 2 | from .module import Module 3 | 4 | 5 | class PairwiseDistance(Module): 6 | norm: float 7 | eps: float 8 | keepdim: bool 9 | 10 | def __init__(self, p: float = ..., eps: float = ..., keepdim: bool = ...) -> None: ... 11 | 12 | def forward(self, x1: Tensor, x2: Tensor) -> Tensor: ... # type: ignore 13 | 14 | def __call__(self, x1: Tensor, x2: Tensor) -> Tensor: ... # type: ignore 15 | 16 | 17 | class CosineSimilarity(Module): 18 | dim: int 19 | eps: float 20 | 21 | def __init__(self, dim: int = ..., eps: float = ...) -> None: ... 22 | 23 | def forward(self, x1: Tensor, x2: Tensor) -> Tensor: ... # type: ignore 24 | 25 | def __call__(self, x1: Tensor, x2: Tensor) -> Tensor: ... # type: ignore 26 | -------------------------------------------------------------------------------- /nn/quantized/modules/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | def _quantize_weight(float_wt, observer): 4 | if observer.qscheme in [torch.per_tensor_symmetric, torch.per_tensor_affine]: 5 | wt_scale, wt_zp = observer.calculate_qparams() 6 | qweight = torch.quantize_per_tensor( 7 | float_wt, 8 | float(wt_scale), int(wt_zp), torch.qint8) 9 | elif observer.qscheme in [torch.per_channel_symmetric, torch.per_channel_affine]: 10 | wt_scale, wt_zp, wt_axis = observer.calculate_qparams() 11 | qweight = torch.quantize_per_channel( 12 | float_wt, 13 | wt_scale.to(torch.double), wt_zp.to(torch.int64), wt_axis, torch.qint8) 14 | else: 15 | raise ValueError("Unexpected qscheme " + observer.qscheme) 16 | return qweight 17 | -------------------------------------------------------------------------------- /nn/modules/flatten.py: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | 3 | class Flatten(Module): 4 | r""" 5 | Flattens a contiguous range of dims into a tensor. For use with :class:`~nn.Sequential`. 6 | Args: 7 | start_dim: first dim to flatten (default = 1). 8 | end_dim: last dim to flatten (default = -1). 9 | 10 | Shape: 11 | - Input: :math:`(N, *dims)` 12 | - Output: :math:`(N, \prod *dims)` (for the default case). 13 | 14 | 15 | Examples:: 16 | >>> m = nn.Sequential( 17 | >>> nn.Conv2d(1, 32, 5, 1, 1), 18 | >>> nn.Flatten() 19 | >>> ) 20 | """ 21 | __constants__ = ['start_dim', 'end_dim'] 22 | 23 | def __init__(self, start_dim=1, end_dim=-1): 24 | super(Flatten, self).__init__() 25 | self.start_dim = start_dim 26 | self.end_dim = end_dim 27 | 28 | def forward(self, input): 29 | return input.flatten(self.start_dim, self.end_dim) 30 | -------------------------------------------------------------------------------- /nn/utils/fusion.py: -------------------------------------------------------------------------------- 1 | 2 | from __future__ import absolute_import, division, print_function, unicode_literals 3 | 4 | import copy 5 | import torch 6 | 7 | def fuse_conv_bn_eval(conv, bn): 8 | assert(not (conv.training or bn.training)), "Fusion only for eval!" 9 | fused_conv = copy.deepcopy(conv) 10 | 11 | fused_conv.weight, fused_conv.bias = \ 12 | fuse_conv_bn_weights(fused_conv.weight, fused_conv.bias, 13 | bn.running_mean, bn.running_var, bn.eps, bn.weight, bn.bias) 14 | 15 | return fused_conv 16 | 17 | def fuse_conv_bn_weights(conv_w, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b): 18 | if conv_b is None: 19 | conv_b = bn_rm.new_zeros(bn_rm.shape) 20 | bn_var_rsqrt = torch.rsqrt(bn_rv + bn_eps) 21 | 22 | conv_w = conv_w * (bn_w * bn_var_rsqrt).reshape([-1] + [1] * (len(conv_w.shape) - 1)) 23 | conv_b = (conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b 24 | 25 | return torch.nn.Parameter(conv_w), torch.nn.Parameter(conv_b) 26 | -------------------------------------------------------------------------------- /nn/parallel/data_parallel.pyi: -------------------------------------------------------------------------------- 1 | from typing import Any, Optional, TypeVar 2 | from .common_types import _devices_t, _device_t 3 | from ..modules import Module 4 | from ... import device, Tensor 5 | 6 | T_co = TypeVar('T_co', covariant=True) 7 | class DataParallel(Module[T_co]): 8 | module: Module = ... 9 | device_ids: _devices_t = ... 10 | dim: int = ... 11 | output_device: _device_t = ... 12 | src_device_obj: device = ... 13 | 14 | def __init__(self, module: Module[T_co], device_ids: Optional[_devices_t] = ..., output_device: Optional[_device_t] = ..., 15 | dim: int = ...) -> None: ... 16 | 17 | def forward(self, *inputs: Any, **kwargs: Any) -> T_co: ... 18 | def __call__(self, *inputs: Any, **kwargs: Any) -> T_co: ... 19 | 20 | 21 | def data_parallel(module: Module, inputs: Any, device_ids: Optional[_devices_t] = ..., 22 | output_device: Optional[_device_t] = ..., dim: int = ..., 23 | module_kwargs: Optional[Any] = ...) -> Tensor: ... 24 | -------------------------------------------------------------------------------- /nn/modules/upsampling.pyi: -------------------------------------------------------------------------------- 1 | from ... import Tensor 2 | from .module import Module 3 | from typing import Optional 4 | from ..common_types import _size_2_t, _ratio_2_t, _size_any_t, _ratio_any_t 5 | 6 | 7 | class Upsample(Module): 8 | name: str = ... 9 | size: _size_any_t = ... 10 | scale_factor: _ratio_any_t = ... 11 | mode: str = ... 12 | align_corners: bool = ... 13 | 14 | def __init__(self, size: Optional[_size_any_t] = ..., scale_factor: Optional[_ratio_any_t] = ..., mode: str = ..., 15 | align_corners: Optional[bool] = ...) -> None: ... 16 | 17 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 18 | 19 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 20 | 21 | 22 | class UpsamplingNearest2d(Upsample): 23 | def __init__(self, size: Optional[_size_2_t] = ..., scale_factor: Optional[_ratio_2_t] = ...) -> None: ... 24 | 25 | 26 | class UpsamplingBilinear2d(Upsample): 27 | def __init__(self, size: Optional[_size_2_t] = ..., scale_factor: Optional[_ratio_2_t] = ...) -> None: ... 28 | -------------------------------------------------------------------------------- /nn/modules/adaptive.pyi: -------------------------------------------------------------------------------- 1 | from ... import Tensor 2 | from .module import Module 3 | from .linear import Linear 4 | from collections import namedtuple 5 | from typing import List, Sequence 6 | from .container import ModuleList 7 | 8 | _ASMoutput = namedtuple('ASMoutput', ['output', 'loss']) 9 | 10 | 11 | class AdaptiveLogSoftmaxWithLoss(Module): 12 | in_features: int = ... 13 | n_classes: int = ... 14 | cutoffs: List[int] = ... 15 | div_value: float = ... 16 | head_bias: bool = ... 17 | head: Linear = ... 18 | tail: ModuleList = ... 19 | 20 | def __init__(self, in_features: int, n_classes: int, cutoffs: Sequence[int], div_value: float = ..., 21 | head_bias: bool = ...) -> None: ... 22 | 23 | def reset_parameters(self) -> None: ... 24 | 25 | def forward(self, input: Tensor, target: Tensor) -> _ASMoutput: ... # type: ignore 26 | 27 | def __call__(self, input: Tensor, target: Tensor) -> _ASMoutput: ... # type: ignore 28 | 29 | def log_prob(self, input: Tensor) -> List[float]: ... 30 | 31 | def predict(self, input: Tensor) -> Tensor: ... 32 | -------------------------------------------------------------------------------- /nn/parallel/scatter_gather.pyi: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, List, Tuple, overload, TypeVar 2 | from ... import Tensor 3 | from .common_types import _device_t, _devices_t 4 | 5 | 6 | T = TypeVar('T', Dict, List, Tuple) 7 | 8 | # For some reason, 'scatter' returns a tuple when given a single Tensor input but a list otherwise. 9 | @overload 10 | def scatter(inputs: Tensor, target_gpus: _devices_t, dim: int = ...) -> Tuple[Tensor, ...]: ... 11 | 12 | # flake8 will raise a spurious error here since `torch/__init__.pyi` has not been generated yet 13 | # so mypy will interpret `Tensor` as `Any` since it is an import from what it belives to be an 14 | # untyped module. Thus to mypy, the first definition of `scatter` looks strictly more general 15 | # than this overload. 16 | @overload 17 | def scatter(inputs: T, target_gpus: _devices_t, dim: int = ...) -> List[T]: ... # type: ignore 18 | 19 | 20 | # TODO More precise types here. 21 | def scatter_kwargs(inputs: Any, kwargs: Any, target_gpus: _devices_t, dim: int = ...) -> Any: ... 22 | 23 | 24 | def gather(outputs: Any, target_device: _device_t, dim: int = ...) -> Any: ... 25 | -------------------------------------------------------------------------------- /nn/parallel/distributed.pyi: -------------------------------------------------------------------------------- 1 | from ..modules import Module 2 | from typing import Any, Optional, TypeVar 3 | from .common_types import _devices_t, _device_t 4 | 5 | T_co = TypeVar('T_co', covariant=True) 6 | 7 | 8 | class DistributedDataParallel(Module[T_co]): 9 | process_group: Any = ... 10 | dim: int = ... 11 | module: Module[T_co] = ... 12 | device_ids: _devices_t = ... 13 | output_device: _device_t = ... 14 | broadcast_buffers: bool = ... 15 | check_reduction: bool = ... 16 | broadcast_bucket_size: float = ... 17 | bucket_bytes_cap: float = ... 18 | 19 | # TODO type process_group once `distributed` module is stubbed 20 | def __init__(self, module: Module[T_co], device_ids: Optional[_devices_t] = ..., 21 | output_device: Optional[_device_t] = ..., dim: int = ..., 22 | broadcast_buffers: bool = ..., process_group: Optional[Any] = ..., bucket_cap_mb: float = ..., 23 | check_reduction: bool = ...) -> None: ... 24 | 25 | def forward(self, *inputs: Any, **kwargs: Any) -> T_co: ... 26 | 27 | def __call__(self, *inputs: Any, **kwargs: Any) -> T_co: ... 28 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # LaLaRAND: Flexible Layer-by-Layer CPU/GPU Scheduling for Real-Time DNN Tasks 2 | 3 | LaLaRAND: Flexible Layer-by-Layer CPU/GPU Scheduling for Real-Time DNN Tasks
4 | Author: Woosung Kang, Kilho Lee, Jinkyu Lee, Insik Shin, Hoon Sung Chwa
5 | In 42nd IEEE Real-Time Systems Symposium (RTSS 2021) Dortmund, Germany, December 2021
6 | 7 | ## Requirements 8 | CUDA: >= 10.2
9 | cuDNN: >=8.0.2
10 | PyTorch: 1.4.0
11 | Python: >= 3.6
12 | CMake: >= 3.10.2
13 | 14 | ## How to use 15 | ### PyTorch modification 16 | 1. Install PyTorch with version 1.4.0 17 | 2. Go to installation directory (probably /home/{username}/.local/lib/python{version}/site-packages/torch} 18 | 3. Replace directory nn, quantization. 19 | 20 | ### Scheduler 21 | 1. Run scheduler before DNN tasks 22 | 2. Provide resource configuration of DNN tasks by txt file (current: /tmp/{pid of task}.txt) 23 | 24 | ### DNN task 25 | 1. Before inference code, 26 | 2. Call {model}.set_rt() to set rt-priority of task 27 | 3. Call {model}.hetero() to use heterogeous resource allocation 28 | 4. hetero() requires inference function and sample inputs for input calibaration 29 | -------------------------------------------------------------------------------- /nn/modules/fold.pyi: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from ... import Tensor 3 | from ..common_types import _size_any_t 4 | 5 | 6 | class Fold(Module): 7 | output_size: _size_any_t = ... 8 | kernel_size: _size_any_t = ... 9 | dilation: _size_any_t = ... 10 | padding: _size_any_t = ... 11 | stride: _size_any_t = ... 12 | 13 | def __init__(self, output_size: _size_any_t, kernel_size: _size_any_t, dilation: _size_any_t = ..., 14 | padding: _size_any_t = ..., stride: _size_any_t = ...) -> None: ... 15 | 16 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 17 | 18 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 19 | 20 | 21 | class Unfold(Module): 22 | kernel_size: _size_any_t = ... 23 | dilation: _size_any_t = ... 24 | padding: _size_any_t = ... 25 | stride: _size_any_t = ... 26 | 27 | def __init__(self, kernel_size: _size_any_t, dilation: _size_any_t = ..., padding: _size_any_t = ..., 28 | stride: _size_any_t = ...) -> None: ... 29 | 30 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 31 | 32 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 33 | -------------------------------------------------------------------------------- /nn/modules/dropout.pyi: -------------------------------------------------------------------------------- 1 | from ... import Tensor 2 | from .module import Module 3 | 4 | 5 | class _DropoutNd(Module): 6 | p: float 7 | inplace: bool 8 | 9 | def __init__(self, p: float = ..., inplace: bool = ...) -> None: ... 10 | 11 | def extra_repr(self): ... 12 | 13 | 14 | class Dropout(_DropoutNd): 15 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 16 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 17 | 18 | 19 | class Dropout2d(_DropoutNd): 20 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 21 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 22 | 23 | 24 | class Dropout3d(_DropoutNd): 25 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 26 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 27 | 28 | 29 | class AlphaDropout(_DropoutNd): 30 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 31 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 32 | 33 | 34 | class FeatureAlphaDropout(_DropoutNd): 35 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 36 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 37 | -------------------------------------------------------------------------------- /scheduler/lalarand.h: -------------------------------------------------------------------------------- 1 | #ifndef _LALARADN_H_ 2 | #define _LALARADN_H_ 3 | 4 | typedef enum{ 5 | IDLE, BUSY 6 | }STATE; 7 | 8 | typedef struct _DNN_INFO{ 9 | int pid; 10 | int id; 11 | int layers; 12 | int period; 13 | double deadline; 14 | int request_fd; 15 | int decision_fd; 16 | int priority; 17 | int current_layer; 18 | int assigned; 19 | int * default_cfg; 20 | struct _DNN_INFO * next; 21 | }dnn_info; 22 | 23 | typedef struct _DNN_QUEUE{ 24 | int count; 25 | dnn_info * head; 26 | }dnn_queue; 27 | 28 | // Queue Definition 29 | typedef struct QNode{ 30 | int layer; 31 | int id; 32 | int priority; 33 | struct QNode * next; 34 | }QNode; 35 | 36 | typedef struct Queue{ 37 | int count; 38 | QNode * front; 39 | QNode * rear; 40 | }Queue; 41 | 42 | typedef struct _MSG_PACKET{ 43 | int regist; 44 | int pid; 45 | int layers; 46 | int period; 47 | int priority; 48 | }reg_msg; 49 | 50 | typedef struct _REQUEST_MSG{ 51 | int request_layer; 52 | }req_msg; 53 | 54 | typedef struct _RESOURCE{ 55 | int res_id; 56 | STATE state; 57 | Queue * waiting; 58 | int id; 59 | int layer; 60 | double scheduled; 61 | }resource; 62 | 63 | #endif 64 | -------------------------------------------------------------------------------- /nn/intrinsic/quantized/modules/linear_relu.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | import torch.nn.quantized as nnq 3 | import torch.nn.intrinsic 4 | import torch 5 | 6 | class LinearReLU(nnq.Linear): 7 | r""" 8 | A LinearReLU module fused from Linear and ReLU modules 9 | 10 | We adopt the same interface as :class:`torch.nn.quantized.Linear`. 11 | 12 | Attributes: 13 | Same as torch.nn.quantized.Linear 14 | 15 | Examples:: 16 | 17 | >>> m = nn.intrinsic.LinearReLU(20, 30) 18 | >>> input = torch.randn(128, 20) 19 | >>> output = m(input) 20 | >>> print(output.size()) 21 | torch.Size([128, 30]) 22 | """ 23 | _FLOAT_MODULE = torch.nn.intrinsic.LinearReLU 24 | 25 | def __init__(self, in_features, out_features, bias=True): 26 | super(LinearReLU, self).__init__(in_features, out_features, bias) 27 | 28 | def forward(self, input): 29 | Y_q = torch.ops.quantized.linear_relu( 30 | input, self._packed_params._packed_params, 31 | float(self.scale), 32 | int(self.zero_point)) 33 | return Y_q 34 | 35 | def _get_name(self): 36 | return 'QuantizedLinearReLU' 37 | 38 | @classmethod 39 | def from_float(cls, mod): 40 | return super(LinearReLU, cls).from_float(mod) 41 | -------------------------------------------------------------------------------- /nn/modules/linear.pyi: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from .. import Parameter 3 | from ... import Tensor 4 | 5 | 6 | class Identity(Module): 7 | 8 | def __init__(self) -> None: ... 9 | 10 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 11 | 12 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 13 | 14 | 15 | class Linear(Module): 16 | in_features: int = ... 17 | out_features: int = ... 18 | weight: Parameter = ... 19 | bias: Parameter = ... 20 | 21 | def __init__(self, in_features: int, out_features: int, bias: bool = ...) -> None: ... 22 | 23 | def reset_parameters(self) -> None: ... 24 | 25 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 26 | 27 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 28 | 29 | 30 | class Bilinear(Module): 31 | in1_features: int = ... 32 | in2_features: int = ... 33 | out_features: int = ... 34 | weight: Parameter = ... 35 | bias: Parameter = ... 36 | 37 | def __init__(self, in1_features: int, in2_features: int, out_features: int, bias: bool = ...) -> None: ... 38 | 39 | def reset_parameters(self) -> None: ... 40 | 41 | def forward(self, input1: Tensor, input2: Tensor) -> Tensor: ... # type: ignore 42 | 43 | def __call__(self, input1: Tensor, input2: Tensor) -> Tensor: ... # type: ignore 44 | -------------------------------------------------------------------------------- /nn/modules/batchnorm.pyi: -------------------------------------------------------------------------------- 1 | from ... import Tensor 2 | from .. import Parameter 3 | from .module import Module 4 | from typing import Any, Optional 5 | 6 | 7 | class _BatchNorm(Module): 8 | num_features: int = ... 9 | eps: float = ... 10 | momentum: float = ... 11 | affine: bool = ... 12 | track_running_stats: bool = ... 13 | weight: Parameter = ... 14 | bias: Parameter = ... 15 | 16 | def __init__(self, num_features: int, eps: float = ..., momentum: float = ..., affine: bool = ..., 17 | track_running_stats: bool = ...) -> None: ... 18 | 19 | def reset_running_stats(self) -> None: ... 20 | 21 | def reset_parameters(self) -> None: ... 22 | 23 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 24 | 25 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 26 | 27 | 28 | class BatchNorm1d(_BatchNorm): ... 29 | 30 | 31 | class BatchNorm2d(_BatchNorm): ... 32 | 33 | 34 | class BatchNorm3d(_BatchNorm): ... 35 | 36 | 37 | class SyncBatchNorm(_BatchNorm): 38 | # TODO set process_group to the write type once torch.distributed is stubbed 39 | def __init__(self, num_features: int, eps: float = ..., momentum: float = ..., affine: bool = ..., 40 | track_running_stats: bool = ..., process_group: Optional[Any] = ...) -> None: ... 41 | 42 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 43 | 44 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 45 | -------------------------------------------------------------------------------- /quantization/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | from .quantize import * 3 | from .observer import * 4 | from .qconfig import * 5 | from .fake_quantize import * 6 | from .fuse_modules import fuse_modules 7 | from .stubs import * 8 | 9 | def default_eval_fn(model, calib_data): 10 | r""" 11 | Default evaluation function takes a torch.utils.data.Dataset or a list of 12 | input Tensors and run the model on the dataset 13 | """ 14 | for data, target in calib_data: 15 | model(data) 16 | 17 | _all__ = [ 18 | 'QuantWrapper', 'QuantStub', 'DeQuantStub', 19 | # Top level API for eager mode quantization 20 | 'quantize', 21 | # Sub functions used by eager mode quantization 22 | 'prepare', 'convert', 23 | # Sub functions for `prepare` and `swap_module` 24 | 'propagate_qconfig_', 'add_quant_dequant', 'add_observer_', 'swap_module', 25 | 'default_eval_fn', 'get_observer_dict', 26 | # Observers 27 | 'ObserverBase', 'WeightObserver', 'observer', 'default_observer', 28 | 'default_weight_observer', 29 | # QConfig 30 | 'QConfig', 'default_qconfig', 'default_dynamic_qconfig', 'float16_dynamic_qconfig', 31 | # QAT utilities 32 | 'default_qat_qconfig', 'prepare_qat', 'quantize_qat', 33 | # module transformations 34 | 'fuse_modules', 35 | # Dynamic quantization utilities 36 | 'quantize_dynamic', 37 | # lalarand layer level quantization 38 | 'prepare_lala' 39 | ] 40 | -------------------------------------------------------------------------------- /nn/intrinsic/qat/modules/linear_relu.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | import torch.nn.qat as nnqat 3 | import torch.nn.intrinsic 4 | import torch.nn.functional as F 5 | 6 | class LinearReLU(nnqat.Linear): 7 | r""" 8 | A LinearReLU module fused from Linear and ReLU modules, attached with 9 | FakeQuantize modules for output activation and weight, used in 10 | quantization aware training. 11 | 12 | We adopt the same interface as :class:`torch.nn.Linear`. 13 | 14 | Similar to `torch.nn.intrinsic.LinearReLU`, with FakeQuantize modules initialized to 15 | default. 16 | 17 | Attributes: 18 | activation_post_process: fake quant module for output activation 19 | weight: fake quant module for weight 20 | 21 | Examples:: 22 | 23 | >>> m = nn.qat.LinearReLU(20, 30) 24 | >>> input = torch.randn(128, 20) 25 | >>> output = m(input) 26 | >>> print(output.size()) 27 | torch.Size([128, 30]) 28 | """ 29 | _FLOAT_MODULE = torch.nn.intrinsic.LinearReLU 30 | 31 | def __init__(self, in_features, out_features, bias=True, 32 | qconfig=None): 33 | super(LinearReLU, self).__init__(in_features, out_features, bias, qconfig) 34 | 35 | def forward(self, input): 36 | return self.activation_post_process(F.relu( 37 | F.linear(input, self.weight_fake_quant(self.weight), self.bias))) 38 | 39 | @classmethod 40 | def from_float(cls, mod, qconfig=None): 41 | return super(LinearReLU, cls).from_float(mod, qconfig) 42 | -------------------------------------------------------------------------------- /nn/common_types.pyi: -------------------------------------------------------------------------------- 1 | from typing import TypeVar, Union, Tuple 2 | from .. import Tensor 3 | 4 | # Create some useful type aliases 5 | 6 | # Template for arguments which can be supplied as a tuple, or which can be a scalar which PyTorch will internally 7 | # broadcast to a tuple. 8 | # Comes in several variants: A tuple of unknown size, and a fixed-size tuple for 1d, 2d, or 3d operations. 9 | T = TypeVar('T') 10 | _scalar_or_tuple_any_t = Union[T, Tuple[T, ...]] 11 | _scalar_or_tuple_1_t = Union[T, Tuple[T]] 12 | _scalar_or_tuple_2_t = Union[T, Tuple[T, T]] 13 | _scalar_or_tuple_3_t = Union[T, Tuple[T, T, T]] 14 | _scalar_or_tuple_4_t = Union[T, Tuple[T, T, T, T]] 15 | _scalar_or_tuple_5_t = Union[T, Tuple[T, T, T, T, T]] 16 | _scalar_or_tuple_6_t = Union[T, Tuple[T, T, T, T, T, T]] 17 | 18 | # For arguments which represent size parameters (eg, kernel size, padding) 19 | _size_any_t = _scalar_or_tuple_any_t[int] 20 | _size_1_t = _scalar_or_tuple_1_t[int] 21 | _size_2_t = _scalar_or_tuple_2_t[int] 22 | _size_3_t = _scalar_or_tuple_3_t[int] 23 | _size_4_t = _scalar_or_tuple_4_t[int] 24 | _size_5_t = _scalar_or_tuple_5_t[int] 25 | _size_6_t = _scalar_or_tuple_6_t[int] 26 | 27 | # For arguments that represent a ratio to adjust each dimension of an input with (eg, upsampling parameters) 28 | _ratio_2_t = _scalar_or_tuple_2_t[float] 29 | _ratio_3_t = _scalar_or_tuple_3_t[float] 30 | _ratio_any_t = _scalar_or_tuple_any_t[float] 31 | 32 | _tensor_list_t = _scalar_or_tuple_any_t[Tensor] 33 | 34 | # For the return value of max pooling operations that may or may not return indices. 35 | # With the proposed 'Literal' feature to Python typing, it might be possible to 36 | # eventually eliminate this. 37 | _maybe_indices_t = _scalar_or_tuple_2_t[Tensor] 38 | -------------------------------------------------------------------------------- /nn/_reduction.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | 3 | # NB: Keep this file in sync with enums in aten/src/ATen/core/Reduction.h 4 | 5 | 6 | def get_enum(reduction): 7 | # type: (str) -> int 8 | if reduction == 'none': 9 | ret = 0 10 | elif reduction == 'mean': 11 | ret = 1 12 | elif reduction == 'elementwise_mean': 13 | warnings.warn("reduction='elementwise_mean' is deprecated, please use reduction='mean' instead.") 14 | ret = 1 15 | elif reduction == 'sum': 16 | ret = 2 17 | else: 18 | ret = -1 # TODO: remove once JIT exceptions support control flow 19 | raise ValueError("{} is not a valid value for reduction".format(reduction)) 20 | return ret 21 | 22 | # In order to support previous versions, accept boolean size_average and reduce 23 | # and convert them into the new constants for now 24 | 25 | 26 | # We use these functions in torch/legacy as well, in which case we'll silence the warning 27 | def legacy_get_string(size_average, reduce, emit_warning=True): 28 | # type: (Optional[bool], Optional[bool], bool) -> str 29 | warning = "size_average and reduce args will be deprecated, please use reduction='{}' instead." 30 | 31 | if size_average is None: 32 | size_average = True 33 | if reduce is None: 34 | reduce = True 35 | 36 | if size_average and reduce: 37 | ret = 'mean' 38 | elif reduce: 39 | ret = 'sum' 40 | else: 41 | ret = 'none' 42 | if emit_warning: 43 | warnings.warn(warning.format(ret)) 44 | return ret 45 | 46 | 47 | def legacy_get_enum(size_average, reduce, emit_warning=True): 48 | # type: (Optional[bool], Optional[bool], bool) -> int 49 | return get_enum(legacy_get_string(size_average, reduce, emit_warning)) 50 | -------------------------------------------------------------------------------- /nn/modules/pixelshuffle.py: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from .. import functional as F 3 | 4 | 5 | class PixelShuffle(Module): 6 | r"""Rearranges elements in a tensor of shape :math:`(*, C \times r^2, H, W)` 7 | to a tensor of shape :math:`(*, C, H \times r, W \times r)`. 8 | 9 | This is useful for implementing efficient sub-pixel convolution 10 | with a stride of :math:`1/r`. 11 | 12 | Look at the paper: 13 | `Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network`_ 14 | by Shi et. al (2016) for more details. 15 | 16 | Args: 17 | upscale_factor (int): factor to increase spatial resolution by 18 | 19 | Shape: 20 | - Input: :math:`(N, L, H_{in}, W_{in})` where :math:`L=C \times \text{upscale\_factor}^2` 21 | - Output: :math:`(N, C, H_{out}, W_{out})` where 22 | :math:`H_{out} = H_{in} \times \text{upscale\_factor}` 23 | and :math:`W_{out} = W_{in} \times \text{upscale\_factor}` 24 | 25 | Examples:: 26 | 27 | >>> pixel_shuffle = nn.PixelShuffle(3) 28 | >>> input = torch.randn(1, 9, 4, 4) 29 | >>> output = pixel_shuffle(input) 30 | >>> print(output.size()) 31 | torch.Size([1, 1, 12, 12]) 32 | 33 | .. _Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network: 34 | https://arxiv.org/abs/1609.05158 35 | """ 36 | __constants__ = ['upscale_factor'] 37 | 38 | def __init__(self, upscale_factor): 39 | super(PixelShuffle, self).__init__() 40 | self.upscale_factor = upscale_factor 41 | 42 | def forward(self, input): 43 | return F.pixel_shuffle(input, self.upscale_factor) 44 | 45 | def extra_repr(self): 46 | return 'upscale_factor={}'.format(self.upscale_factor) 47 | -------------------------------------------------------------------------------- /nn/parameter.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from collections import OrderedDict 3 | 4 | 5 | class Parameter(torch.Tensor): 6 | r"""A kind of Tensor that is to be considered a module parameter. 7 | 8 | Parameters are :class:`~torch.Tensor` subclasses, that have a 9 | very special property when used with :class:`Module` s - when they're 10 | assigned as Module attributes they are automatically added to the list of 11 | its parameters, and will appear e.g. in :meth:`~Module.parameters` iterator. 12 | Assigning a Tensor doesn't have such effect. This is because one might 13 | want to cache some temporary state, like last hidden state of the RNN, in 14 | the model. If there was no such class as :class:`Parameter`, these 15 | temporaries would get registered too. 16 | 17 | Arguments: 18 | data (Tensor): parameter tensor. 19 | requires_grad (bool, optional): if the parameter requires gradient. See 20 | :ref:`excluding-subgraphs` for more details. Default: `True` 21 | """ 22 | 23 | def __new__(cls, data=None, requires_grad=True): 24 | if data is None: 25 | data = torch.Tensor() 26 | return torch.Tensor._make_subclass(cls, data, requires_grad) 27 | 28 | def __deepcopy__(self, memo): 29 | if id(self) in memo: 30 | return memo[id(self)] 31 | else: 32 | result = type(self)(self.data.clone(memory_format=torch.preserve_format), self.requires_grad) 33 | memo[id(self)] = result 34 | return result 35 | 36 | def __repr__(self): 37 | return 'Parameter containing:\n' + super(Parameter, self).__repr__() 38 | 39 | def __reduce_ex__(self, proto): 40 | # See Note [Don't serialize hooks] 41 | return ( 42 | torch._utils._rebuild_parameter, 43 | (self.data, self.requires_grad, OrderedDict()) 44 | ) 45 | -------------------------------------------------------------------------------- /quantization/stubs.py: -------------------------------------------------------------------------------- 1 | 2 | from torch import nn 3 | 4 | class QuantStub(nn.Module): 5 | r"""Quantize stub module, before calibration, this is same as an observer, 6 | it will be swapped as `nnq.Quantize` in `convert`. 7 | 8 | Args: 9 | qconfig: quantization configuration for the tensor, 10 | if qconfig is not provided, we will get qconfig from parent modules 11 | """ 12 | def __init__(self, qconfig=None): 13 | super(QuantStub, self).__init__() 14 | if qconfig: 15 | self.qconfig = qconfig 16 | 17 | def forward(self, x): 18 | return x 19 | 20 | 21 | class DeQuantStub(nn.Module): 22 | r"""Dequantize stub module, before calibration, this is same as identity, 23 | this will be swapped as `nnq.DeQuantize` in `convert`. 24 | """ 25 | def __init__(self): 26 | super(DeQuantStub, self).__init__() 27 | 28 | def forward(self, x): 29 | return x 30 | 31 | 32 | class QuantWrapper(nn.Module): 33 | r"""A wrapper class that wraps the input module, adds QuantStub and 34 | DeQuantStub and surround the call to module with call to quant and dequant 35 | modules. 36 | 37 | This is used by the `quantization` utility functions to add the quant and 38 | dequant modules, before `convert` function `QuantStub` will just be observer, 39 | it observes the input tensor, after `convert`, `QuantStub` 40 | will be swapped to `nnq.Quantize` which does actual quantization. Similarly 41 | for `DeQuantStub`. 42 | """ 43 | def __init__(self, module): 44 | super(QuantWrapper, self).__init__() 45 | qconfig = module.qconfig if hasattr(module, 'qconfig') else None 46 | self.add_module('quant', QuantStub(qconfig)) 47 | self.add_module('dequant', DeQuantStub()) 48 | self.add_module('module', module) 49 | self.train(module.training) 50 | 51 | def forward(self, X): 52 | X = self.quant(X) 53 | X = self.module(X) 54 | return self.dequant(X) 55 | -------------------------------------------------------------------------------- /scheduler/lalarand_fn.h: -------------------------------------------------------------------------------- 1 | #ifndef _LALALRAND_FN_H_ 2 | #define _LALARAND_FN_H_ 3 | #include "lalarand.h" 4 | 5 | void set_priority(int priority); 6 | void set_affinity(int core); 7 | 8 | ///// dnn queue //// 9 | dnn_queue * createDNNQueue(); 10 | void enDNNQueue(dnn_queue * dnn_list, dnn_info* dnn); 11 | void deleteDNN(dnn_queue * dnn_list, dnn_info* dnn); 12 | void setDNNpriority(dnn_queue * dnn_list); 13 | 14 | ///// resource ///// 15 | resource * createResource(int res_id); 16 | 17 | //// waiting queue //// 18 | QNode* newNode (int layer, int id, int priority); 19 | Queue * createQueue(); 20 | void enQueue(Queue *q, int layer, int id, int priority); 21 | int deQueue(Queue * q, double current_time, resource * res); 22 | dnn_info * find_dnn_by_id(dnn_queue * dnn_list, int id); 23 | dnn_info * find_dnn_by_pid(dnn_queue * dnn_list, int pid); 24 | int find_node_by_id(Queue *q, int id); 25 | void print_queue(char * name, Queue * q); 26 | void print_list(char * name, dnn_queue * dnn_list); 27 | void del_arg(int argc, char **argv, int index); 28 | int find_int_arg(int argc, char **argv, char *arg, int def); 29 | double get_time_point(); 30 | void send_release_time(dnn_queue * dnn_list); 31 | void check_registration(dnn_queue * dnn_list, int reg_fd, resource * gpu, resource *cpu); 32 | void regist(dnn_queue * dnn_list, reg_msg * msg); 33 | void de_regist(dnn_queue * dnn_list, reg_msg *msg, resource * gpu, resource * cpu); 34 | void request_handler(dnn_info * node, resource * gpu, resource * cpu, double current_time); 35 | void decision_handler(int target_id, dnn_queue * dnn_list, int decision); 36 | void update_deadline(dnn_info * dnn, double current_time); 37 | void update_deadline_all(dnn_queue * dnn_list, double current_time); 38 | char* get_resource_name(int id); 39 | int make_fdset(fd_set *readfds,int reg_fd, dnn_queue * dnn_list); 40 | int open_channel(char * pipe_name,int mode); 41 | void close_channel(char * pipe_name); 42 | void close_channels(dnn_info * dnn); 43 | void read_default_cfg(int pid, int * default_cfg); 44 | #endif 45 | -------------------------------------------------------------------------------- /nn/modules/normalization.pyi: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from typing import Any, Union, List 3 | from ... import Tensor, Size 4 | from .. import Parameter 5 | 6 | 7 | class LocalResponseNorm(Module): 8 | size: int = ... 9 | alpha: float = ... 10 | beta: float = ... 11 | k: float = ... 12 | 13 | def __init__(self, size: int, alpha: float = ..., beta: float = ..., k: float = ...) -> None: ... 14 | 15 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 16 | 17 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 18 | 19 | 20 | class CrossMapLRN2d(Module): 21 | size: int = ... 22 | alpha: float = ... 23 | beta: float = ... 24 | k: float = ... 25 | 26 | def __init__(self, size: int, alpha: float = ..., beta: float = ..., k: float = ...) -> None: ... 27 | 28 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 29 | 30 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 31 | 32 | 33 | _shape_t = Union[int, List[int], Size] 34 | 35 | 36 | class LayerNorm(Module): 37 | normalized_shape: _shape_t = ... 38 | eps: float = ... 39 | elementwise_affine: bool = ... 40 | weight: Parameter = ... 41 | bias: Parameter = ... 42 | 43 | def __init__(self, normalized_shape: _shape_t, eps: float = ..., elementwise_affine: bool = ...) -> None: ... 44 | 45 | def reset_parameters(self) -> None: ... 46 | 47 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 48 | 49 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 50 | 51 | 52 | class GroupNorm(Module): 53 | num_groups: int = ... 54 | num_channels: int = ... 55 | eps: float = ... 56 | affine: bool = ... 57 | weight: Parameter = ... 58 | bias: Parameter = ... 59 | 60 | def __init__(self, num_groups: int, num_channels: int, eps: float = ..., affine: bool = ...) -> None: ... 61 | 62 | def reset_parameters(self) -> None: ... 63 | 64 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 65 | 66 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 67 | -------------------------------------------------------------------------------- /quantization/default_mappings.py: -------------------------------------------------------------------------------- 1 | 2 | from torch import nn 3 | 4 | import torch.nn.intrinsic as nni 5 | import torch.nn.intrinsic.quantized as nniq 6 | import torch.nn.intrinsic.qat as nniqat 7 | import torch.nn.quantized as nnq 8 | import torch.nn.quantized.dynamic as nnqd 9 | import torch.nn.qat as nnqat 10 | 11 | from .stubs import QuantStub, DeQuantStub 12 | 13 | # Map for swapping float module to quantized ones 14 | DEFAULT_MODULE_MAPPING = { 15 | nn.Linear: nnq.Linear, 16 | nn.ReLU: nnq.ReLU, 17 | nn.ReLU6: nnq.ReLU6, 18 | nn.Conv2d: nnq.Conv2d, 19 | nn.Conv3d: nnq.Conv3d, 20 | QuantStub: nnq.Quantize, 21 | DeQuantStub: nnq.DeQuantize, 22 | # Wrapper Modules: 23 | nnq.FloatFunctional: nnq.QFunctional, 24 | # Intrinsic modules: 25 | nni.ConvReLU2d: nniq.ConvReLU2d, 26 | nni.ConvReLU3d: nniq.ConvReLU3d, 27 | nni.LinearReLU: nniq.LinearReLU, 28 | nniqat.ConvReLU2d: nniq.ConvReLU2d, 29 | nniqat.LinearReLU: nniq.LinearReLU, 30 | nniqat.ConvBn2d: nnq.Conv2d, 31 | nniqat.ConvBnReLU2d: nniq.ConvReLU2d, 32 | # QAT modules: 33 | nnqat.Linear: nnq.Linear, 34 | nnqat.Conv2d: nnq.Conv2d, 35 | } 36 | 37 | # Map for swapping float module to qat modules 38 | DEFAULT_QAT_MODULE_MAPPING = { 39 | nn.Linear: nnqat.Linear, 40 | nn.Conv2d: nnqat.Conv2d, 41 | # Intrinsic modules: 42 | nni.ConvBn2d: nniqat.ConvBn2d, 43 | nni.ConvBnReLU2d: nniqat.ConvBnReLU2d, 44 | nni.ConvReLU2d: nniqat.ConvReLU2d, 45 | nni.LinearReLU: nniqat.LinearReLU 46 | } 47 | 48 | # Map for swapping dynamic modules 49 | DEFAULT_DYNAMIC_MODULE_MAPPING = { 50 | nn.Linear: nnqd.Linear, 51 | nn.LSTM: nnqd.LSTM, 52 | } 53 | 54 | # Whitelist for propagating the qconfig 55 | _EXCLUDE_QCONFIG_PROPAGATE_LIST = { 56 | DeQuantStub, 57 | } 58 | _INCLUDE_QCONFIG_PROPAGATE_LIST = { 59 | nn.Sequential, 60 | nn.MaxPool2d, 61 | nn.AdaptiveAvgPool2d, 62 | } 63 | 64 | DEFAULT_QCONFIG_PROPAGATE_WHITE_LIST = ( 65 | set(DEFAULT_MODULE_MAPPING.keys()) | 66 | set(DEFAULT_QAT_MODULE_MAPPING.keys()) | 67 | set(DEFAULT_DYNAMIC_MODULE_MAPPING.keys()) | 68 | _INCLUDE_QCONFIG_PROPAGATE_LIST - 69 | _EXCLUDE_QCONFIG_PROPAGATE_LIST 70 | ) 71 | -------------------------------------------------------------------------------- /nn/modules/padding.pyi: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from ... import Tensor 3 | from ..common_types import _size_2_t, _size_4_t, _size_6_t 4 | 5 | 6 | class _ConstantPadNd(Module): 7 | value: float 8 | 9 | def __init__(self, value: float) -> None: ... 10 | 11 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 12 | 13 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 14 | 15 | 16 | class ConstantPad1d(_ConstantPadNd): 17 | padding: _size_2_t = ... 18 | 19 | def __init__(self, padding: _size_2_t, value: float) -> None: ... 20 | 21 | 22 | class ConstantPad2d(_ConstantPadNd): 23 | padding: _size_4_t = ... 24 | 25 | def __init__(self, padding: _size_4_t, value: float) -> None: ... 26 | 27 | 28 | class ConstantPad3d(_ConstantPadNd): 29 | padding: _size_6_t = ... 30 | 31 | def __init__(self, padding: _size_6_t, value: float) -> None: ... 32 | 33 | 34 | class _ReflectionPadNd(Module): 35 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 36 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 37 | 38 | def extra_repr(self): ... 39 | 40 | 41 | class ReflectionPad1d(_ReflectionPadNd): 42 | padding: _size_2_t = ... 43 | 44 | def __init__(self, padding: _size_2_t) -> None: ... 45 | 46 | 47 | class ReflectionPad2d(_ReflectionPadNd): 48 | padding: _size_4_t = ... 49 | 50 | def __init__(self, padding: _size_4_t) -> None: ... 51 | 52 | 53 | class _ReplicationPadNd(Module): 54 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 55 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 56 | 57 | def extra_repr(self): ... 58 | 59 | 60 | class ReplicationPad1d(_ReplicationPadNd): 61 | padding: _size_2_t = ... 62 | 63 | def __init__(self, padding: _size_2_t) -> None: ... 64 | 65 | 66 | class ReplicationPad2d(_ReplicationPadNd): 67 | padding: _size_4_t = ... 68 | 69 | def __init__(self, padding: _size_4_t) -> None: ... 70 | 71 | 72 | class ReplicationPad3d(_ReplicationPadNd): 73 | padding: _size_6_t = ... 74 | 75 | def __init__(self, padding: _size_6_t) -> None: ... 76 | 77 | 78 | class ZeroPad2d(ConstantPad2d): 79 | padding: _size_4_t = ... 80 | 81 | def __init__(self, padding: _size_4_t) -> None: ... 82 | -------------------------------------------------------------------------------- /nn/modules/sparse.pyi: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from typing import Optional 3 | from .. import Parameter 4 | from ... import Tensor 5 | 6 | 7 | class Embedding(Module): 8 | num_embeddings: int = ... 9 | embedding_dim: int = ... 10 | padding_idx: int = ... 11 | max_norm: float = ... 12 | norm_type: float = ... 13 | scale_grad_by_freq: bool = ... 14 | weight: Parameter = ... 15 | sparse: bool = ... 16 | 17 | def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = ..., 18 | max_norm: Optional[float] = ..., norm_type: float = ..., scale_grad_by_freq: bool = ..., 19 | sparse: bool = ..., _weight: Optional[Tensor] = ...) -> None: ... 20 | 21 | def reset_parameters(self) -> None: ... 22 | 23 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 24 | 25 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 26 | 27 | @classmethod 28 | def from_pretrained(cls, embeddings: Tensor, freeze: bool = ..., padding_idx: Optional[int] = ..., 29 | max_norm: Optional[float] = ..., norm_type: float = ..., scale_grad_by_freq: bool = ..., 30 | sparse: bool = ...): ... 31 | 32 | 33 | class EmbeddingBag(Module): 34 | num_embeddings: int = ... 35 | embedding_dim: int = ... 36 | max_norm: float = ... 37 | norm_type: float = ... 38 | scale_grad_by_freq: bool = ... 39 | weight: Parameter = ... 40 | mode: str = ... 41 | sparse: bool = ... 42 | 43 | def __init__(self, num_embeddings: int, embedding_dim: int, max_norm: Optional[float] = ..., norm_type: float = ..., 44 | scale_grad_by_freq: bool = ..., mode: str = ..., sparse: bool = ..., 45 | _weight: Optional[Tensor] = ...) -> None: ... 46 | 47 | def reset_parameters(self) -> None: ... 48 | 49 | def forward(self, input: Tensor, offsets: Optional[Tensor] = ...) -> Tensor: ... # type: ignore 50 | 51 | def __call__(self, input: Tensor, offsets: Optional[Tensor] = ...) -> Tensor: ... # type: ignore 52 | 53 | @classmethod 54 | def from_pretrained(cls, embeddings: Tensor, freeze: bool = ..., max_norm: Optional[float] = ..., 55 | norm_type: float = ..., scale_grad_by_freq: bool = ..., mode: str = ..., 56 | sparse: bool = ...): ... 57 | -------------------------------------------------------------------------------- /nn/qat/modules/linear.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn.intrinsic import LinearReLU 5 | 6 | class Linear(nn.Linear): 7 | r""" 8 | A linear module attached with FakeQuantize modules for both output 9 | activation and weight, used for quantization aware training. 10 | 11 | We adopt the same interface as `torch.nn.Linear`, please see 12 | https://pytorch.org/docs/stable/nn.html#torch.nn.Linear 13 | for documentation. 14 | 15 | Similar to `torch.nn.Linear`, with FakeQuantize modules initialized to 16 | default. 17 | 18 | Attributes: 19 | activation_post_process: fake quant module for output activation 20 | weight: fake quant module for weight 21 | """ 22 | _FLOAT_MODULE = nn.Linear 23 | 24 | def __init__(self, in_features, out_features, bias=True, 25 | qconfig=None): 26 | super(Linear, self).__init__(in_features, out_features, bias) 27 | assert qconfig, 'qconfig must be provided for QAT module' 28 | self.qconfig = qconfig 29 | self.activation_post_process = qconfig.activation() 30 | self.weight_fake_quant = qconfig.weight() 31 | 32 | def forward(self, input): 33 | return self.activation_post_process( 34 | F.linear(input, self.weight_fake_quant(self.weight), self.bias)) 35 | 36 | @classmethod 37 | def from_float(cls, mod, qconfig=None): 38 | r"""Create a qat module from a float module or qparams_dict 39 | 40 | Args: `mod` a float module, either produced by torch.quantization utilities 41 | or directly from user 42 | """ 43 | assert type(mod) == cls._FLOAT_MODULE, ' qat.' + cls.__name__ + '.from_float only works for ' + \ 44 | cls._FLOAT_MODULE.__name__ 45 | if not qconfig: 46 | assert hasattr(mod, 'qconfig'), 'Input float module must have qconfig defined' 47 | assert mod.qconfig, 'Input float module must have a valid qconfig' 48 | if type(mod) == LinearReLU: 49 | mod = mod[0] 50 | 51 | qconfig = mod.qconfig 52 | qat_linear = cls(mod.in_features, mod.out_features, bias=mod.bias is not None, qconfig=qconfig) 53 | qat_linear.weight = mod.weight 54 | qat_linear.bias = mod.bias 55 | return qat_linear 56 | -------------------------------------------------------------------------------- /nn/intrinsic/modules/fused.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | import torch 3 | from torch.nn import Conv2d, Conv3d, ReLU, Linear, BatchNorm2d 4 | 5 | class ConvReLU2d(torch.nn.Sequential): 6 | r"""This is a sequential container which calls the Conv 2d and ReLU modules. 7 | During quantization this will be replaced with the corresponding fused module.""" 8 | def __init__(self, conv, relu): 9 | assert type(conv) == Conv2d and type(relu) == ReLU, \ 10 | 'Incorrect types for input modules{}{}'.format( 11 | type(conv), type(relu)) 12 | super(ConvReLU2d, self).__init__(conv, relu) 13 | 14 | class ConvReLU3d(torch.nn.Sequential): 15 | r"""This is a sequential container which calls the Conv 3d and ReLU modules. 16 | During quantization this will be replaced with the corresponding fused module.""" 17 | def __init__(self, conv, relu): 18 | assert type(conv) == Conv3d and type(relu) == ReLU, \ 19 | 'Incorrect types for input modules{}{}'.format( 20 | type(conv), type(relu)) 21 | super(ConvReLU3d, self).__init__(conv, relu) 22 | 23 | class LinearReLU(torch.nn.Sequential): 24 | r"""This is a sequential container which calls the Linear and ReLU modules. 25 | During quantization this will be replaced with the corresponding fused module.""" 26 | def __init__(self, linear, relu): 27 | assert type(linear) == Linear and type(relu) == ReLU, \ 28 | 'Incorrect types for input modules{}{}'.format( 29 | type(linear), type(relu)) 30 | super(LinearReLU, self).__init__(linear, relu) 31 | 32 | class ConvBn2d(torch.nn.Sequential): 33 | r"""This is a sequential container which calls the Conv 2d and Batch Norm 2d modules. 34 | During quantization this will be replaced with the corresponding fused module.""" 35 | def __init__(self, conv, bn): 36 | assert type(conv) == Conv2d and type(bn) == BatchNorm2d, \ 37 | 'Incorrect types for input modules{}{}'.format( 38 | type(conv), type(bn)) 39 | super(ConvBn2d, self).__init__(conv, bn) 40 | 41 | class ConvBnReLU2d(torch.nn.Sequential): 42 | r"""This is a sequential container which calls the Conv 2d, Batch Norm 2d, and ReLU modules. 43 | During quantization this will be replaced with the corresponding fused module.""" 44 | def __init__(self, conv, bn, relu): 45 | assert type(conv) == Conv2d and type(bn) == BatchNorm2d and \ 46 | type(relu) == ReLU, 'Incorrect types for input modules{}{}{}' \ 47 | .format(type(conv), type(bn), type(relu)) 48 | super(ConvBnReLU2d, self).__init__(conv, bn, relu) 49 | -------------------------------------------------------------------------------- /nn/modules/distance.py: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from .. import functional as F 3 | 4 | 5 | class PairwiseDistance(Module): 6 | r""" 7 | Computes the batchwise pairwise distance between vectors :math:`v_1`, :math:`v_2` using the p-norm: 8 | 9 | .. math :: 10 | \Vert x \Vert _p = \left( \sum_{i=1}^n \vert x_i \vert ^ p \right) ^ {1/p}. 11 | 12 | Args: 13 | p (real): the norm degree. Default: 2 14 | eps (float, optional): Small value to avoid division by zero. 15 | Default: 1e-6 16 | keepdim (bool, optional): Determines whether or not to keep the vector dimension. 17 | Default: False 18 | Shape: 19 | - Input1: :math:`(N, D)` where `D = vector dimension` 20 | - Input2: :math:`(N, D)`, same shape as the Input1 21 | - Output: :math:`(N)`. If :attr:`keepdim` is ``True``, then :math:`(N, 1)`. 22 | Examples:: 23 | >>> pdist = nn.PairwiseDistance(p=2) 24 | >>> input1 = torch.randn(100, 128) 25 | >>> input2 = torch.randn(100, 128) 26 | >>> output = pdist(input1, input2) 27 | """ 28 | __constants__ = ['norm', 'eps', 'keepdim'] 29 | 30 | def __init__(self, p=2., eps=1e-6, keepdim=False): 31 | super(PairwiseDistance, self).__init__() 32 | self.norm = p 33 | self.eps = eps 34 | self.keepdim = keepdim 35 | 36 | def forward(self, x1, x2): 37 | return F.pairwise_distance(x1, x2, self.norm, self.eps, self.keepdim) 38 | 39 | 40 | class CosineSimilarity(Module): 41 | r"""Returns cosine similarity between :math:`x_1` and :math:`x_2`, computed along dim. 42 | 43 | .. math :: 44 | \text{similarity} = \dfrac{x_1 \cdot x_2}{\max(\Vert x_1 \Vert _2 \cdot \Vert x_2 \Vert _2, \epsilon)}. 45 | 46 | Args: 47 | dim (int, optional): Dimension where cosine similarity is computed. Default: 1 48 | eps (float, optional): Small value to avoid division by zero. 49 | Default: 1e-8 50 | Shape: 51 | - Input1: :math:`(\ast_1, D, \ast_2)` where D is at position `dim` 52 | - Input2: :math:`(\ast_1, D, \ast_2)`, same shape as the Input1 53 | - Output: :math:`(\ast_1, \ast_2)` 54 | Examples:: 55 | >>> input1 = torch.randn(100, 128) 56 | >>> input2 = torch.randn(100, 128) 57 | >>> cos = nn.CosineSimilarity(dim=1, eps=1e-6) 58 | >>> output = cos(input1, input2) 59 | """ 60 | __constants__ = ['dim', 'eps'] 61 | 62 | def __init__(self, dim=1, eps=1e-8): 63 | super(CosineSimilarity, self).__init__() 64 | self.dim = dim 65 | self.eps = eps 66 | 67 | def forward(self, x1, x2): 68 | return F.cosine_similarity(x1, x2, self.dim, self.eps) 69 | -------------------------------------------------------------------------------- /scheduler/lalarand.c: -------------------------------------------------------------------------------- 1 | #define DEBUG 0 2 | #ifndef _GNU_SOURCE 3 | #define _GNU_SOURCE 4 | #endif 5 | 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | #include 19 | #include 20 | #include 21 | 22 | #include "lalarand.h" 23 | #include "lalarand_fn.h" 24 | 25 | #define GPU 1 26 | #define CPU 0 27 | 28 | #define REGISTRATION "/tmp/lalarand_registration" 29 | 30 | 31 | int main(int argc, char **argv){ 32 | int Sync = find_int_arg(argc, argv, "-sync", 1); 33 | 34 | set_priority(50); 35 | set_affinity(7); 36 | 37 | dnn_queue * dnn_list = createDNNQueue(); 38 | resource * gpu = createResource(GPU); 39 | resource * cpu = createResource(CPU); 40 | 41 | int reg_fd = open_channel(REGISTRATION, O_RDONLY | O_NONBLOCK); 42 | 43 | double current_time; 44 | int gpu_target, cpu_target; 45 | int fd_head; 46 | fd_set readfds; 47 | dnn_info *node; 48 | 49 | do{ 50 | gpu_target = -1; 51 | cpu_target = -1; 52 | 53 | fd_head = make_fdset(&readfds, reg_fd, dnn_list); 54 | 55 | if(select(fd_head +1, &readfds, NULL, NULL, NULL)){ 56 | current_time = get_time_point(); 57 | if(FD_ISSET(reg_fd, &readfds)) { 58 | check_registration(dnn_list, reg_fd, gpu, cpu); 59 | print_list("REGIST",dnn_list); 60 | } 61 | for(node = dnn_list ->head; node !=NULL; node = node -> next){ 62 | if(FD_ISSET(node->request_fd, &readfds)){ 63 | request_handler(node, gpu, cpu, current_time); 64 | } 65 | } 66 | 67 | print_queue("GPU",gpu->waiting); 68 | print_queue("CPU",cpu->waiting); 69 | 70 | if(!(gpu->waiting->count + cpu->waiting->count < Sync)){ 71 | if(Sync) update_deadline_all(dnn_list, current_time); 72 | 73 | if(gpu -> state == IDLE) gpu_target = deQueue(gpu->waiting, current_time, gpu); 74 | if(cpu -> state == IDLE) cpu_target = deQueue(cpu->waiting, current_time, cpu); 75 | 76 | if(Sync) send_release_time(dnn_list); 77 | 78 | if(gpu_target != -1) decision_handler(gpu_target, dnn_list, GPU); 79 | if(cpu_target != -1) decision_handler(cpu_target, dnn_list, CPU); 80 | Sync = 0; 81 | } 82 | } 83 | }while(!(Sync == 0 && dnn_list -> count == 0)); 84 | } 85 | -------------------------------------------------------------------------------- /nn/quantized/modules/activation.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | from __future__ import unicode_literals 5 | 6 | import torch 7 | import torch.nn.quantized.functional 8 | 9 | class ReLU(torch.nn.ReLU): 10 | r"""Applies quantized rectified linear unit function element-wise: 11 | 12 | :math:`\text{ReLU}(x)= \max(x_0, x)`, where :math:`x_0` is the zero point. 13 | 14 | Please see https://pytorch.org/docs/stable/nn.html#torch.nn.ReLU 15 | for more documentation on ReLU. 16 | 17 | Args: 18 | inplace: (Currently not supported) can optionally do the operation in-place. 19 | 20 | Shape: 21 | - Input: :math:`(N, *)` where `*` means, any number of additional 22 | dimensions 23 | - Output: :math:`(N, *)`, same shape as the input 24 | 25 | Examples:: 26 | 27 | >>> m = nn.quantized.ReLU() 28 | >>> input = torch.randn(2) 29 | >>> input = torch.quantize_per_tensor(input, 1.0, 0, dtype=torch.qint32) 30 | >>> output = m(input) 31 | """ 32 | def __init__(self, inplace=False): 33 | super(ReLU, self).__init__(inplace) 34 | self.inplace = inplace 35 | 36 | def forward(self, input): 37 | return torch.nn.quantized.functional.relu(input, inplace=self.inplace) 38 | 39 | def _get_name(self): 40 | return 'QuantizedReLU' 41 | 42 | @staticmethod 43 | def from_float(mod): 44 | return ReLU(mod.inplace) 45 | 46 | 47 | class ReLU6(torch.nn.ReLU): 48 | r"""Applies the element-wise function: 49 | 50 | :math:`\text{ReLU6}(x) = \min(\max(x_0, x), q(6))`, where :math:`x_0` is the 51 | zero_point, and :math:`q(6)` is the quantized representation of number 6. 52 | 53 | Args: 54 | inplace: can optionally do the operation in-place. Default: ``False`` 55 | 56 | Shape: 57 | - Input: :math:`(N, *)` where `*` means, any number of additional 58 | dimensions 59 | - Output: :math:`(N, *)`, same shape as the input 60 | 61 | .. image:: scripts/activation_images/ReLU6.png 62 | 63 | Examples:: 64 | 65 | >>> m = nn.quantized.ReLU6() 66 | >>> input = torch.randn(2) 67 | >>> input = torch.quantize_per_tensor(input, 1.0, 0, dtype=torch.qint32) 68 | >>> output = m(input) 69 | """ 70 | def __init__(self, inplace=False): 71 | super(ReLU6, self).__init__(inplace) 72 | self.inplace = inplace 73 | 74 | def forward(self, input): 75 | return torch.ops.quantized.relu6(input, self.inplace) 76 | 77 | def _get_name(self): 78 | return 'QuantizedReLU6' 79 | 80 | @staticmethod 81 | def from_float(mod): 82 | return ReLU6(mod.inplace) 83 | -------------------------------------------------------------------------------- /nn/qat/modules/conv.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | from torch.nn import Conv2d as NNConv2d 3 | from torch.nn.intrinsic import ConvReLU2d 4 | 5 | class Conv2d(NNConv2d): 6 | r""" 7 | A Conv2d module attached with FakeQuantize modules for both output 8 | activation and weight, used for quantization aware training. 9 | 10 | We adopt the same interface as `torch.nn.Conv2d`, please see 11 | https://pytorch.org/docs/stable/nn.html?highlight=conv2d#torch.nn.Conv2d 12 | for documentation. 13 | 14 | Similar to `torch.nn.Conv2d`, with FakeQuantize modules initialized to 15 | default. 16 | 17 | Attributes: 18 | activation_post_process: fake quant module for output activation 19 | weight_fake_quant: fake quant module for weight 20 | """ 21 | _FLOAT_MODULE = NNConv2d 22 | 23 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, 24 | padding=0, dilation=1, groups=1, 25 | bias=True, padding_mode='zeros', qconfig=None): 26 | super(Conv2d, self).__init__(in_channels, out_channels, kernel_size, 27 | stride=stride, padding=padding, dilation=dilation, 28 | groups=groups, bias=bias, padding_mode=padding_mode) 29 | assert qconfig, 'qconfig must be provided for QAT module' 30 | self.qconfig = qconfig 31 | self.activation_post_process = qconfig.activation() 32 | self.weight_fake_quant = qconfig.weight() 33 | 34 | def forward(self, input): 35 | return self.activation_post_process( 36 | self.conv2d_forward(input, self.weight_fake_quant(self.weight))) 37 | 38 | @classmethod 39 | def from_float(cls, mod, qconfig=None): 40 | r"""Create a qat module from a float module or qparams_dict 41 | 42 | Args: `mod` a float module, either produced by torch.quantization utilities 43 | or directly from user 44 | """ 45 | assert type(mod) == cls._FLOAT_MODULE, 'qat.' + cls.__name__ + '.from_float only works for ' + \ 46 | cls._FLOAT_MODULE.__name__ 47 | if not qconfig: 48 | assert hasattr(mod, 'qconfig'), 'Input float module must have qconfig defined' 49 | assert mod.qconfig, 'Input float module must have a valid qconfig' 50 | if type(mod) == ConvReLU2d: 51 | mod = mod[0] 52 | qconfig = mod.qconfig 53 | qat_conv = cls(mod.in_channels, mod.out_channels, mod.kernel_size, 54 | stride=mod.stride, padding=mod.padding, dilation=mod.dilation, 55 | groups=mod.groups, bias=mod.bias is not None, 56 | padding_mode=mod.padding_mode, qconfig=qconfig) 57 | qat_conv.weight = mod.weight 58 | qat_conv.bias = mod.bias 59 | return qat_conv 60 | -------------------------------------------------------------------------------- /nn/utils/clip_grad.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | import torch 3 | from torch._six import inf 4 | 5 | 6 | def clip_grad_norm_(parameters, max_norm, norm_type=2): 7 | r"""Clips gradient norm of an iterable of parameters. 8 | 9 | The norm is computed over all gradients together, as if they were 10 | concatenated into a single vector. Gradients are modified in-place. 11 | 12 | Arguments: 13 | parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a 14 | single Tensor that will have gradients normalized 15 | max_norm (float or int): max norm of the gradients 16 | norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for 17 | infinity norm. 18 | 19 | Returns: 20 | Total norm of the parameters (viewed as a single vector). 21 | """ 22 | if isinstance(parameters, torch.Tensor): 23 | parameters = [parameters] 24 | parameters = list(filter(lambda p: p.grad is not None, parameters)) 25 | max_norm = float(max_norm) 26 | norm_type = float(norm_type) 27 | if norm_type == inf: 28 | total_norm = max(p.grad.data.abs().max() for p in parameters) 29 | else: 30 | total_norm = 0 31 | for p in parameters: 32 | param_norm = p.grad.data.norm(norm_type) 33 | total_norm += param_norm.item() ** norm_type 34 | total_norm = total_norm ** (1. / norm_type) 35 | clip_coef = max_norm / (total_norm + 1e-6) 36 | if clip_coef < 1: 37 | for p in parameters: 38 | p.grad.data.mul_(clip_coef) 39 | return total_norm 40 | 41 | 42 | def clip_grad_norm(parameters, max_norm, norm_type=2): 43 | r"""Clips gradient norm of an iterable of parameters. 44 | 45 | .. warning:: 46 | This method is now deprecated in favor of 47 | :func:`torch.nn.utils.clip_grad_norm_`. 48 | """ 49 | warnings.warn("torch.nn.utils.clip_grad_norm is now deprecated in favor " 50 | "of torch.nn.utils.clip_grad_norm_.", stacklevel=2) 51 | return clip_grad_norm_(parameters, max_norm, norm_type) 52 | 53 | 54 | def clip_grad_value_(parameters, clip_value): 55 | r"""Clips gradient of an iterable of parameters at specified value. 56 | 57 | Gradients are modified in-place. 58 | 59 | Arguments: 60 | parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a 61 | single Tensor that will have gradients normalized 62 | clip_value (float or int): maximum allowed value of the gradients. 63 | The gradients are clipped in the range 64 | :math:`\left[\text{-clip\_value}, \text{clip\_value}\right]` 65 | """ 66 | if isinstance(parameters, torch.Tensor): 67 | parameters = [parameters] 68 | clip_value = float(clip_value) 69 | for p in filter(lambda p: p.grad is not None, parameters): 70 | p.grad.data.clamp_(min=-clip_value, max=clip_value) 71 | -------------------------------------------------------------------------------- /nn/parallel/scatter_gather.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from ._functions import Scatter, Gather 3 | 4 | 5 | def scatter(inputs, target_gpus, dim=0): 6 | r""" 7 | Slices tensors into approximately equal chunks and 8 | distributes them across given GPUs. Duplicates 9 | references to objects that are not tensors. 10 | """ 11 | def scatter_map(obj): 12 | if isinstance(obj, torch.Tensor): 13 | return Scatter.apply(target_gpus, None, dim, obj) 14 | if isinstance(obj, tuple) and len(obj) > 0: 15 | return list(zip(*map(scatter_map, obj))) 16 | if isinstance(obj, list) and len(obj) > 0: 17 | return list(map(list, zip(*map(scatter_map, obj)))) 18 | if isinstance(obj, dict) and len(obj) > 0: 19 | return list(map(type(obj), zip(*map(scatter_map, obj.items())))) 20 | return [obj for targets in target_gpus] 21 | 22 | # After scatter_map is called, a scatter_map cell will exist. This cell 23 | # has a reference to the actual function scatter_map, which has references 24 | # to a closure that has a reference to the scatter_map cell (because the 25 | # fn is recursive). To avoid this reference cycle, we set the function to 26 | # None, clearing the cell 27 | try: 28 | res = scatter_map(inputs) 29 | finally: 30 | scatter_map = None 31 | return res 32 | 33 | 34 | def scatter_kwargs(inputs, kwargs, target_gpus, dim=0): 35 | r"""Scatter with support for kwargs dictionary""" 36 | inputs = scatter(inputs, target_gpus, dim) if inputs else [] 37 | kwargs = scatter(kwargs, target_gpus, dim) if kwargs else [] 38 | if len(inputs) < len(kwargs): 39 | inputs.extend([() for _ in range(len(kwargs) - len(inputs))]) 40 | elif len(kwargs) < len(inputs): 41 | kwargs.extend([{} for _ in range(len(inputs) - len(kwargs))]) 42 | inputs = tuple(inputs) 43 | kwargs = tuple(kwargs) 44 | return inputs, kwargs 45 | 46 | 47 | def gather(outputs, target_device, dim=0): 48 | r""" 49 | Gathers tensors from different GPUs on a specified device 50 | (-1 means the CPU). 51 | """ 52 | def gather_map(outputs): 53 | out = outputs[0] 54 | if isinstance(out, torch.Tensor): 55 | return Gather.apply(target_device, dim, *outputs) 56 | if out is None: 57 | return None 58 | if isinstance(out, dict): 59 | if not all((len(out) == len(d) for d in outputs)): 60 | raise ValueError('All dicts must have the same number of keys') 61 | return type(out)(((k, gather_map([d[k] for d in outputs])) 62 | for k in out)) 63 | return type(out)(map(gather_map, zip(*outputs))) 64 | 65 | # Recursive function calls like this create reference cycles. 66 | # Setting the function to None clears the refcycle. 67 | try: 68 | res = gather_map(outputs) 69 | finally: 70 | gather_map = None 71 | return res 72 | -------------------------------------------------------------------------------- /nn/quantized/modules/__init__.py: -------------------------------------------------------------------------------- 1 | # @lint-ignore-every PYTHON3COMPATIMPORTS 2 | 3 | import torch 4 | from torch.nn.modules.pooling import MaxPool2d 5 | 6 | from .activation import ReLU, ReLU6 7 | from .conv import Conv2d, Conv3d 8 | from .linear import Linear 9 | 10 | from .functional_modules import FloatFunctional, QFunctional 11 | 12 | 13 | class Quantize(torch.nn.Module): 14 | r"""Quantizes an incoming tensor 15 | 16 | Args: 17 | `scale`: scale of the output Quantized Tensor 18 | `zero_point`: zero_point of output Quantized Tensor 19 | `dtype`: data type of output Quantized Tensor 20 | 21 | Attributes: 22 | `scale`, `zero_point`, `dtype` 23 | 24 | Examples:: 25 | >>> t = torch.tensor([[1., -1.], [1., -1.]]) 26 | >>> scale, zero_point, dtype = 1.0, 2, torch.qint8 27 | >>> qm = Quantize(scale, zero_point, dtype) 28 | >>> qt = qm(t) 29 | >>> print(qt) 30 | tensor([[ 1., -1.], 31 | [ 1., -1.]], size=(2, 2), dtype=torch.qint8, scale=1.0, zero_point=2) 32 | """ 33 | 34 | def __init__(self, scale, zero_point, dtype): 35 | super(Quantize, self).__init__() 36 | self.register_buffer('scale', torch.tensor([scale])) 37 | self.register_buffer('zero_point', torch.tensor([zero_point], dtype=torch.long)) 38 | self.dtype = dtype 39 | 40 | def forward(self, X): 41 | return torch.quantize_per_tensor(X, float(self.scale), 42 | int(self.zero_point), self.dtype) 43 | 44 | @staticmethod 45 | def from_float(mod): 46 | assert hasattr(mod, 'activation_post_process') 47 | scale, zero_point = mod.activation_post_process.calculate_qparams() 48 | return Quantize(scale.float().item(), zero_point.long().item(), mod.activation_post_process.dtype) 49 | 50 | def extra_repr(self): 51 | return 'scale={}, zero_point={}, dtype={}'.format(self.scale, self.zero_point, self.dtype) 52 | 53 | 54 | class DeQuantize(torch.nn.Module): 55 | r"""Dequantizes an incoming tensor 56 | 57 | Examples:: 58 | >>> input = torch.tensor([[1., -1.], [1., -1.]]) 59 | >>> scale, zero_point, dtype = 1.0, 2, torch.qint8 60 | >>> qm = Quantize(scale, zero_point, dtype) 61 | >>> quantized_input = qm(input) 62 | >>> dqm = DeQuantize() 63 | >>> dequantized = dqm(quantized_input) 64 | >>> print(dequantized) 65 | tensor([[ 1., -1.], 66 | [ 1., -1.]], dtype=torch.float32) 67 | """ 68 | 69 | def __init__(self): 70 | super(DeQuantize, self).__init__() 71 | 72 | def forward(self, Xq): 73 | return Xq.dequantize() 74 | 75 | @staticmethod 76 | def from_float(mod): 77 | return DeQuantize() 78 | 79 | __all__ = [ 80 | 'Conv2d', 81 | 'Conv3d', 82 | 'DeQuantize', 83 | 'Linear', 84 | 'MaxPool2d', 85 | 'Quantize', 86 | 'ReLU', 87 | 'ReLU6', 88 | # Wrapper modules 89 | 'FloatFunctional', 90 | 'QFunctional', 91 | ] 92 | -------------------------------------------------------------------------------- /nn/cpp.py: -------------------------------------------------------------------------------- 1 | """Functionality for Python <-> C++ frontend inter-op.""" 2 | 3 | from torch import nn 4 | 5 | 6 | class OrderedDictWrapper(object): 7 | """ 8 | A wrapper around a C++ OrderedDict that dynamically evaluates the 9 | OrderedDict getter on a bound C++ module, such that new changes on the C++ 10 | side are picked up. Otherwise accessing e.g. ``cpp_module._parameters`` just 11 | once would get a frozen copy of the parameters at the time of access. 12 | ``torch.nn.Module`` accesses ``_parameters`` et al. via ``self.__dict__`` so 13 | using properties does not work. 14 | """ 15 | 16 | def __init__(self, cpp_module, attr): 17 | self.cpp_module = cpp_module 18 | self.attr = attr 19 | 20 | @property 21 | def cpp_dict(self): 22 | return getattr(self.cpp_module, self.attr) 23 | 24 | # Magic methods cannot be assigned dynamically and bypass ``getattr``, so we 25 | # must manually override them. 26 | 27 | def items(self): 28 | return self.cpp_dict.items() 29 | 30 | def keys(self): 31 | return self.cpp_dict.keys() 32 | 33 | def values(self): 34 | return self.cpp_dict.values() 35 | 36 | def __iter__(self): 37 | return self.cpp_dict.__iter__() 38 | 39 | def __len__(self): 40 | return self.cpp_dict.__len__() 41 | 42 | def __contains__(self, key): 43 | return self.cpp_dict.__contains__(key) 44 | 45 | def __getitem__(self, key): 46 | return self.cpp_dict.__getitem__(key) 47 | 48 | 49 | class ModuleWrapper(nn.Module): 50 | """ 51 | A subclass of ``torch.nn.Module`` that wraps a C++ frontend module and 52 | delegates all access. 53 | """ 54 | 55 | def __init__(self, cpp_module): 56 | # Assign before the super class constructor so ``self.training`` can be 57 | # assigned to in the super class constructor. 58 | self.cpp_module = cpp_module 59 | super(ModuleWrapper, self).__init__() 60 | self._parameters = OrderedDictWrapper(cpp_module, "_parameters") 61 | self._buffers = OrderedDictWrapper(cpp_module, "_buffers") 62 | self._modules = OrderedDictWrapper(cpp_module, "_modules") 63 | for attr in dir(cpp_module): 64 | # Skip magic methods and the three attributes above. 65 | if not attr.startswith("_"): 66 | setattr(self, attr, getattr(self.cpp_module, attr)) 67 | 68 | def _apply(self, fn): 69 | for param in self.parameters(): 70 | # Tensors stored in modules are graph leaves, and we don't 71 | # want to create copy nodes, so we have to unpack the data. 72 | param.data = fn(param.data) 73 | if param._grad is not None: 74 | param._grad.data = fn(param._grad.data) 75 | 76 | for buf in self.buffers(): 77 | buf.data = fn(buf.data) 78 | 79 | return self 80 | 81 | @property 82 | def training(self): 83 | return self.cpp_module.training 84 | 85 | @training.setter 86 | def training(self, mode): 87 | self.cpp_module.train(mode) 88 | 89 | def __repr__(self): 90 | return self.cpp_module.__repr__() 91 | -------------------------------------------------------------------------------- /nn/utils/convert_parameters.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def parameters_to_vector(parameters): 5 | r"""Convert parameters to one vector 6 | 7 | Arguments: 8 | parameters (Iterable[Tensor]): an iterator of Tensors that are the 9 | parameters of a model. 10 | 11 | Returns: 12 | The parameters represented by a single vector 13 | """ 14 | # Flag for the device where the parameter is located 15 | param_device = None 16 | 17 | vec = [] 18 | for param in parameters: 19 | # Ensure the parameters are located in the same device 20 | param_device = _check_param_device(param, param_device) 21 | 22 | vec.append(param.view(-1)) 23 | return torch.cat(vec) 24 | 25 | 26 | def vector_to_parameters(vec, parameters): 27 | r"""Convert one vector to the parameters 28 | 29 | Arguments: 30 | vec (Tensor): a single vector represents the parameters of a model. 31 | parameters (Iterable[Tensor]): an iterator of Tensors that are the 32 | parameters of a model. 33 | """ 34 | # Ensure vec of type Tensor 35 | if not isinstance(vec, torch.Tensor): 36 | raise TypeError('expected torch.Tensor, but got: {}' 37 | .format(torch.typename(vec))) 38 | # Flag for the device where the parameter is located 39 | param_device = None 40 | 41 | # Pointer for slicing the vector for each parameter 42 | pointer = 0 43 | for param in parameters: 44 | # Ensure the parameters are located in the same device 45 | param_device = _check_param_device(param, param_device) 46 | 47 | # The length of the parameter 48 | num_param = param.numel() 49 | # Slice the vector, reshape it, and replace the old data of the parameter 50 | param.data = vec[pointer:pointer + num_param].view_as(param).data 51 | 52 | # Increment the pointer 53 | pointer += num_param 54 | 55 | 56 | def _check_param_device(param, old_param_device): 57 | r"""This helper function is to check if the parameters are located 58 | in the same device. Currently, the conversion between model parameters 59 | and single vector form is not supported for multiple allocations, 60 | e.g. parameters in different GPUs, or mixture of CPU/GPU. 61 | 62 | Arguments: 63 | param ([Tensor]): a Tensor of a parameter of a model 64 | old_param_device (int): the device where the first parameter of a 65 | model is allocated. 66 | 67 | Returns: 68 | old_param_device (int): report device for the first time 69 | """ 70 | 71 | # Meet the first parameter 72 | if old_param_device is None: 73 | old_param_device = param.get_device() if param.is_cuda else -1 74 | else: 75 | warn = False 76 | if param.is_cuda: # Check if in same GPU 77 | warn = (param.get_device() != old_param_device) 78 | else: # Check if in CPU 79 | warn = (old_param_device != -1) 80 | if warn: 81 | raise TypeError('Found two parameters on different devices, ' 82 | 'this is currently not supported.') 83 | return old_param_device 84 | -------------------------------------------------------------------------------- /nn/parallel/parallel_apply.py: -------------------------------------------------------------------------------- 1 | import threading 2 | import torch 3 | from torch.cuda._utils import _get_device_index 4 | from torch._utils import ExceptionWrapper 5 | 6 | 7 | def get_a_var(obj): 8 | if isinstance(obj, torch.Tensor): 9 | return obj 10 | 11 | if isinstance(obj, list) or isinstance(obj, tuple): 12 | for result in map(get_a_var, obj): 13 | if isinstance(result, torch.Tensor): 14 | return result 15 | if isinstance(obj, dict): 16 | for result in map(get_a_var, obj.items()): 17 | if isinstance(result, torch.Tensor): 18 | return result 19 | return None 20 | 21 | 22 | def parallel_apply(modules, inputs, kwargs_tup=None, devices=None): 23 | r"""Applies each `module` in :attr:`modules` in parallel on arguments 24 | contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword) 25 | on each of :attr:`devices`. 26 | 27 | Args: 28 | modules (Module): modules to be parallelized 29 | inputs (tensor): inputs to the modules 30 | devices (list of int or torch.device): CUDA devices 31 | 32 | :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and 33 | :attr:`devices` (if given) should all have same length. Moreover, each 34 | element of :attr:`inputs` can either be a single object as the only argument 35 | to a module, or a collection of positional arguments. 36 | """ 37 | assert len(modules) == len(inputs) 38 | if kwargs_tup is not None: 39 | assert len(modules) == len(kwargs_tup) 40 | else: 41 | kwargs_tup = ({},) * len(modules) 42 | if devices is not None: 43 | assert len(modules) == len(devices) 44 | else: 45 | devices = [None] * len(modules) 46 | devices = list(map(lambda x: _get_device_index(x, True), devices)) 47 | lock = threading.Lock() 48 | results = {} 49 | grad_enabled = torch.is_grad_enabled() 50 | 51 | def _worker(i, module, input, kwargs, device=None): 52 | torch.set_grad_enabled(grad_enabled) 53 | if device is None: 54 | device = get_a_var(input).get_device() 55 | try: 56 | with torch.cuda.device(device): 57 | # this also avoids accidental slicing of `input` if it is a Tensor 58 | if not isinstance(input, (list, tuple)): 59 | input = (input,) 60 | output = module(*input, **kwargs) 61 | with lock: 62 | results[i] = output 63 | except Exception: 64 | with lock: 65 | results[i] = ExceptionWrapper( 66 | where="in replica {} on device {}".format(i, device)) 67 | 68 | if len(modules) > 1: 69 | threads = [threading.Thread(target=_worker, 70 | args=(i, module, input, kwargs, device)) 71 | for i, (module, input, kwargs, device) in 72 | enumerate(zip(modules, inputs, kwargs_tup, devices))] 73 | 74 | for thread in threads: 75 | thread.start() 76 | for thread in threads: 77 | thread.join() 78 | else: 79 | _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0]) 80 | 81 | outputs = [] 82 | for i in range(len(inputs)): 83 | output = results[i] 84 | if isinstance(output, ExceptionWrapper): 85 | output.reraise() 86 | outputs.append(output) 87 | return outputs 88 | -------------------------------------------------------------------------------- /nn/intrinsic/quantized/modules/conv_relu.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | 3 | import torch 4 | import torch.nn.intrinsic 5 | import torch.nn.intrinsic.qat 6 | import torch.nn.quantized as nnq 7 | 8 | from torch.nn.utils import fuse_conv_bn_weights 9 | 10 | 11 | class ConvReLU2d(nnq.Conv2d): 12 | r""" 13 | A ConvReLU2d module is a fused module of Conv2d and ReLU 14 | 15 | We adopt the same interface as :class:`torch.nn.quantized.Conv2d`. 16 | 17 | Attributes: 18 | Same as torch.nn.quantized.Conv2d 19 | 20 | """ 21 | _FLOAT_MODULE = torch.nn.intrinsic.ConvReLU2d 22 | 23 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, 24 | padding=0, dilation=1, groups=1, bias=True, 25 | padding_mode='zeros'): 26 | super(ConvReLU2d, self).__init__( 27 | in_channels, out_channels, kernel_size, stride=stride, 28 | padding=padding, dilation=dilation, groups=groups, bias=bias, 29 | padding_mode=padding_mode) 30 | 31 | def forward(self, input): 32 | # Temporarily using len(shape) instead of ndim due to JIT issue 33 | # https://github.com/pytorch/pytorch/issues/23890 34 | if len(input.shape) != 4: 35 | raise ValueError("Input shape must be `(N, C, H, W)`!") 36 | return torch.ops.quantized.conv2d_relu( 37 | input, self._packed_params, self.stride, self.padding, 38 | self.dilation, self.groups, self.scale, self.zero_point) 39 | 40 | def _get_name(self): 41 | return 'QuantizedConvReLU2d' 42 | 43 | @classmethod 44 | def from_float(cls, mod): 45 | if type(mod) == torch.nn.intrinsic.qat.ConvBnReLU2d: 46 | mod.weight, mod.bias = fuse_conv_bn_weights( 47 | mod.weight, mod.bias, mod.running_mean, mod.running_var, 48 | mod.eps, mod.gamma, mod.beta) 49 | return super(ConvReLU2d, cls).from_float(mod) 50 | 51 | 52 | class ConvReLU3d(nnq.Conv3d): 53 | r""" 54 | A ConvReLU3d module is a fused module of Conv3d and ReLU 55 | 56 | We adopt the same interface as :class:`torch.nn.quantized.Conv3d`. 57 | 58 | .. note:: 59 | Attributes: Same as torch.nn.quantized.Conv3d 60 | 61 | """ 62 | _FLOAT_MODULE = torch.nn.intrinsic.ConvReLU3d 63 | 64 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, 65 | padding=0, dilation=1, groups=1, bias=True, 66 | padding_mode='zeros'): 67 | super(ConvReLU3d, self).__init__( 68 | in_channels, out_channels, kernel_size, stride=stride, 69 | padding=padding, dilation=dilation, groups=groups, bias=bias, 70 | padding_mode=padding_mode) 71 | 72 | def forward(self, input): 73 | # Temporarily using len(shape) instead of ndim due to JIT issue 74 | # https://github.com/pytorch/pytorch/issues/23890 75 | if len(input.shape) != 5: 76 | raise ValueError("Input shape must be `(N, C, D, H, W)`!") 77 | return torch.ops.quantized.conv3d_relu( 78 | input, self._packed_params, self.stride, self.padding, 79 | self.dilation, self.groups, self.scale, self.zero_point) 80 | 81 | def _get_name(self): 82 | return 'QuantizedConvReLU3d' 83 | 84 | @classmethod 85 | def from_float(cls, mod): 86 | if type(mod) == torch.nn.intrinsic.qat.ConvBnReLU3d: 87 | mod.weight, mod.bias = fuse_conv_bn_weights( 88 | mod.weight, mod.bias, mod.running_mean, mod.running_var, 89 | mod.eps, mod.gamma, mod.beta) 90 | return super(ConvReLU3d, cls).from_float(mod) 91 | -------------------------------------------------------------------------------- /nn/quantized/dynamic/modules/linear.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | import torch 3 | from ....modules.linear import Linear as NNLinear 4 | import torch.nn.quantized as nnq 5 | from torch.nn.quantized.modules.utils import _quantize_weight 6 | 7 | class Linear(nnq.Linear): 8 | r""" 9 | A dynamic quantized linear module with quantized tensor as inputs and outputs. 10 | We adopt the same interface as `torch.nn.Linear`, please see 11 | https://pytorch.org/docs/stable/nn.html#torch.nn.Linear for documentation. 12 | 13 | Similar to :class:`torch.nn.Linear`, attributes will be randomly 14 | initialized at module creation time and will be overwritten later 15 | 16 | Attributes: 17 | weight (Tensor): the non-learnable quantized weights of the module which are of 18 | shape :math:`(\text{out\_features}, \text{in\_features})`. 19 | bias (Tensor): the non-learnable bias of the module of shape :math:`(\text{out\_features})`. 20 | If :attr:`bias` is ``True``, the values are initialized to zero. 21 | 22 | Examples:: 23 | 24 | >>> m = nn.quantized.dynamic.Linear(20, 30) 25 | >>> input = torch.randn(128, 20) 26 | >>> output = m(input) 27 | >>> print(output.size()) 28 | torch.Size([128, 30]) 29 | """ 30 | 31 | def __init__(self, in_features, out_features, bias_=True): 32 | super(Linear, self).__init__(in_features, out_features, bias_) 33 | # We don't muck around with buffers or attributes or anything here 34 | # to keep the module simple. *everything* is simply a Python attribute. 35 | # Serialization logic is explicitly handled in the below serialization and 36 | # deserialization modules 37 | 38 | def forward(self, x): 39 | # Note that we can handle self.bias == None case. 40 | Y = torch.ops.quantized.linear_dynamic( 41 | x, self._packed_params._packed_params) 42 | return Y.to(x.dtype) 43 | 44 | def _get_name(self): 45 | return 'DynamicQuantizedLinear' 46 | 47 | def extra_repr(self): 48 | return 'in_features={}, out_features={}'.format( 49 | self.in_features, self.out_features 50 | ) 51 | 52 | @classmethod 53 | def from_float(cls, mod): 54 | r"""Create a dynamic quantized module from a float module or qparams_dict 55 | 56 | Args: 57 | mod (Module): a float module, either produced by torch.quantization 58 | utilities or provided by the user 59 | """ 60 | assert type(mod) == NNLinear, 'nn.quantized.dynamic.Linear.from_float only works for nn.Linear' 61 | assert hasattr(mod, 'qconfig'), 'Input float module must have qconfig defined' 62 | if mod.qconfig is not None and mod.qconfig.weight is not None: 63 | weight_observer = mod.qconfig.weight() 64 | else: 65 | # We have the circular import issues if we import the qconfig in the beginning of this file: 66 | # https://github.com/pytorch/pytorch/pull/24231. The current workaround is to postpone the 67 | # import until we need it. 68 | from torch.quantization.qconfig import default_dynamic_qconfig 69 | weight_observer = default_dynamic_qconfig.weight() 70 | assert weight_observer.dtype == torch.qint8, 'Weight observer must have dtype torch.qint8' 71 | weight_observer(mod.weight) 72 | qweight = _quantize_weight(mod.weight.float(), weight_observer) 73 | qlinear = Linear(mod.in_features, mod.out_features) 74 | qlinear.set_weight_bias(qweight, mod.bias) 75 | return qlinear 76 | -------------------------------------------------------------------------------- /nn/modules/__init__.pyi: -------------------------------------------------------------------------------- 1 | from .module import Module as Module 2 | from .activation import CELU as CELU, ELU as ELU, GLU as GLU, GELU as GELU, Hardshrink as Hardshrink, \ 3 | Hardtanh as Hardtanh, LeakyReLU as LeakyReLU, LogSigmoid as LogSigmoid, LogSoftmax as LogSoftmax, PReLU as PReLU, \ 4 | RReLU as RReLU, ReLU as ReLU, ReLU6 as ReLU6, SELU as SELU, Sigmoid as Sigmoid, Softmax as Softmax, \ 5 | Softmax2d as Softmax2d, Softmin as Softmin, Softplus as Softplus, Softshrink as Softshrink, Softsign as Softsign, \ 6 | Tanh as Tanh, Tanhshrink as Tanhshrink, Threshold as Threshold 7 | from .adaptive import AdaptiveLogSoftmaxWithLoss as AdaptiveLogSoftmaxWithLoss 8 | from .batchnorm import BatchNorm1d as BatchNorm1d, BatchNorm2d as BatchNorm2d, BatchNorm3d as BatchNorm3d, \ 9 | SyncBatchNorm as SyncBatchNorm 10 | from .container import Container as Container, ModuleDict as ModuleDict, ModuleList as ModuleList, \ 11 | ParameterDict as ParameterDict, ParameterList as ParameterList, Sequential as Sequential 12 | from .conv import Conv1d as Conv1d, Conv2d as Conv2d, Conv3d as Conv3d, ConvTranspose1d as ConvTranspose1d, \ 13 | ConvTranspose2d as ConvTranspose2d, ConvTranspose3d as ConvTranspose3d 14 | from .distance import CosineSimilarity as CosineSimilarity, PairwiseDistance as PairwiseDistance 15 | from .dropout import AlphaDropout as AlphaDropout, Dropout as Dropout, Dropout2d as Dropout2d, Dropout3d as Dropout3d, \ 16 | FeatureAlphaDropout as FeatureAlphaDropout 17 | from .fold import Fold as Fold, Unfold as Unfold 18 | from .instancenorm import InstanceNorm1d as InstanceNorm1d, InstanceNorm2d as InstanceNorm2d, \ 19 | InstanceNorm3d as InstanceNorm3d 20 | from .linear import Bilinear as Bilinear, Identity as Identity, Linear as Linear 21 | from .loss import BCELoss as BCELoss, BCEWithLogitsLoss as BCEWithLogitsLoss, CTCLoss as CTCLoss, \ 22 | CosineEmbeddingLoss as CosineEmbeddingLoss, CrossEntropyLoss as CrossEntropyLoss, \ 23 | HingeEmbeddingLoss as HingeEmbeddingLoss, KLDivLoss as KLDivLoss, L1Loss as L1Loss, MSELoss as MSELoss, \ 24 | MarginRankingLoss as MarginRankingLoss, MultiLabelMarginLoss as MultiLabelMarginLoss, \ 25 | MultiLabelSoftMarginLoss as MultiLabelSoftMarginLoss, MultiMarginLoss as MultiMarginLoss, NLLLoss as NLLLoss, \ 26 | NLLLoss2d as NLLLoss2d, PoissonNLLLoss as PoissonNLLLoss, SmoothL1Loss as SmoothL1Loss, \ 27 | SoftMarginLoss as SoftMarginLoss, TripletMarginLoss as TripletMarginLoss 28 | from .module import Module as Module 29 | from .normalization import CrossMapLRN2d as CrossMapLRN2d, GroupNorm as GroupNorm, LayerNorm as LayerNorm, \ 30 | LocalResponseNorm as LocalResponseNorm 31 | from .padding import ConstantPad1d as ConstantPad1d, ConstantPad2d as ConstantPad2d, ConstantPad3d as ConstantPad3d, \ 32 | ReflectionPad1d as ReflectionPad1d, ReflectionPad2d as ReflectionPad2d, ReplicationPad1d as ReplicationPad1d, \ 33 | ReplicationPad2d as ReplicationPad2d, ReplicationPad3d as ReplicationPad3d, ZeroPad2d as ZeroPad2d 34 | from .pixelshuffle import PixelShuffle as PixelShuffle 35 | from .pooling import AdaptiveAvgPool1d as AdaptiveAvgPool1d, AdaptiveAvgPool2d as AdaptiveAvgPool2d, \ 36 | AdaptiveAvgPool3d as AdaptiveAvgPool3d, AdaptiveMaxPool1d as AdaptiveMaxPool1d, \ 37 | AdaptiveMaxPool2d as AdaptiveMaxPool2d, AdaptiveMaxPool3d as AdaptiveMaxPool3d, AvgPool1d as AvgPool1d, \ 38 | AvgPool2d as AvgPool2d, AvgPool3d as AvgPool3d, FractionalMaxPool2d as FractionalMaxPool2d, \ 39 | FractionalMaxPool3d as FractionalMaxPool3d, LPPool1d as LPPool1d, LPPool2d as LPPool2d, MaxPool1d as MaxPool1d, \ 40 | MaxPool2d as MaxPool2d, MaxPool3d as MaxPool3d, MaxUnpool1d as MaxUnpool1d, MaxUnpool2d as MaxUnpool2d, \ 41 | MaxUnpool3d as MaxUnpool3d 42 | from .rnn import GRU as GRU, GRUCell as GRUCell, LSTM as LSTM, LSTMCell as LSTMCell, RNN as RNN, RNNBase as RNNBase, \ 43 | RNNCell as RNNCell, RNNCellBase as RNNCellBase 44 | from .sparse import Embedding as Embedding, EmbeddingBag as EmbeddingBag 45 | from .upsampling import Upsample as Upsample, UpsamplingBilinear2d as UpsamplingBilinear2d, \ 46 | UpsamplingNearest2d as UpsamplingNearest2d 47 | -------------------------------------------------------------------------------- /nn/modules/container.pyi: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from typing import Any, Optional, Union, overload, TypeVar, Iterable, Tuple, Mapping, Iterator 3 | from collections import OrderedDict 4 | from ... import Tensor 5 | from .. import Parameter 6 | 7 | 8 | class Container(Module): 9 | def __init__(self, **kwargs: Any) -> None: ... 10 | 11 | 12 | T = TypeVar('T') 13 | 14 | 15 | class Sequential(Module): 16 | @overload 17 | def __init__(self, *args: Module) -> None: ... 18 | 19 | @overload 20 | def __init__(self, arg: OrderedDict[str, Module]) -> None: ... 21 | 22 | @overload 23 | def __getitem__(self, idx: int) -> Module: ... 24 | 25 | @overload 26 | def __getitem__(self: T, idx: slice) -> T: ... 27 | 28 | def __setitem__(self, idx: Union[int], module: Module) -> None: ... 29 | 30 | def __delitem__(self, idx: Union[slice, int]) -> None: ... 31 | 32 | def __len__(self) -> int: ... 33 | 34 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 35 | 36 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 37 | 38 | 39 | class ModuleList(Module): 40 | def __init__(self, modules: Optional[Iterable[Module]] = ...) -> None: ... 41 | 42 | @overload 43 | def __getitem__(self, idx: int) -> Module: ... 44 | 45 | @overload 46 | def __getitem__(self: T, idx: slice) -> T: ... 47 | 48 | def __setitem__(self, idx: int, module: Module) -> None: ... 49 | 50 | def __delitem__(self, idx: Union[int, slice]) -> None: ... 51 | 52 | def __len__(self) -> int: ... 53 | 54 | def __iter__(self) -> Iterator[Module]: ... 55 | 56 | def __iadd__(self: T, modules: Iterable[Module]) -> T: ... 57 | 58 | def insert(self, index: int, module: Module) -> None: ... 59 | 60 | def append(self: T, module: Module) -> T: ... 61 | 62 | def extend(self: T, modules: Iterable[Module]) -> T: ... 63 | 64 | 65 | class ModuleDict(Module): 66 | def __init__(self, modules: Optional[Mapping[str, Module]] = ...) -> None: ... 67 | 68 | def __getitem__(self, key: str): ... 69 | 70 | def __setitem__(self, key: str, module: Module) -> None: ... 71 | 72 | def __delitem__(self, key: str) -> None: ... 73 | 74 | def __len__(self) -> int: ... 75 | 76 | def __iter__(self) -> Iterator[str]: ... 77 | 78 | def __contains__(self, key: str) -> bool: ... 79 | 80 | def clear(self) -> None: ... 81 | 82 | def pop(self, key: str): ... 83 | 84 | def keys(self) -> Iterable[str]: ... 85 | 86 | def items(self) -> Iterable[Tuple[str, Module]]: ... 87 | 88 | def values(self) -> Iterable[Module]: ... 89 | 90 | def update(self, modules: Mapping[str, Module]) -> None: ... 91 | 92 | 93 | class ParameterList(Module): 94 | def __init__(self, parameters: Optional[Iterable[Parameter]] = ...) -> None: ... 95 | 96 | @overload 97 | def __getitem__(self, idx: int) -> Parameter: ... 98 | 99 | @overload 100 | def __getitem__(self: T, idx: slice) -> T: ... 101 | 102 | def __setitem__(self, idx: int, param: Parameter) -> None: ... 103 | 104 | def __delitem__(self, idx: Union[int, slice]) -> None: ... 105 | 106 | def __len__(self) -> int: ... 107 | 108 | def __iter__(self) -> Iterator[Parameter]: ... 109 | 110 | def __iadd__(self: T, parameters: Iterable[Parameter]) -> T: ... 111 | 112 | def insert(self, index: int, parameter: Parameter) -> None: ... 113 | 114 | def append(self: T, parameter: Parameter) -> T: ... 115 | 116 | def extend(self: T, parameters: Iterable[Parameter]) -> T: ... 117 | 118 | 119 | class ParameterDict(Module): 120 | def __init__(self, parameters: Optional[Mapping[str, Parameter]] = ...) -> None: ... 121 | 122 | def __getitem__(self, key: str): ... 123 | 124 | def __setitem__(self, key: str, param: Parameter) -> None: ... 125 | 126 | def __delitem__(self, key: str) -> None: ... 127 | 128 | def __len__(self) -> int: ... 129 | 130 | def __iter__(self) -> Iterator[str]: ... 131 | 132 | def __contains__(self, key: str) -> bool: ... 133 | 134 | def clear(self) -> None: ... 135 | 136 | def pop(self, key: str): ... 137 | 138 | def keys(self) -> Iterable[str]: ... 139 | 140 | def items(self) -> Iterable[Tuple[str, Parameter]]: ... 141 | 142 | def values(self) -> Iterable[Parameter]: ... 143 | 144 | def update(self, parameters: Mapping[str, Parameter]) -> None: ... 145 | -------------------------------------------------------------------------------- /nn/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from .linear import Identity, Linear, Bilinear 3 | from .conv import Conv1d, Conv2d, Conv3d, \ 4 | ConvTranspose1d, ConvTranspose2d, ConvTranspose3d 5 | from .activation import Threshold, ReLU, Hardtanh, ReLU6, Sigmoid, Tanh, \ 6 | Softmax, Softmax2d, LogSoftmax, ELU, SELU, CELU, GELU, Hardshrink, LeakyReLU, LogSigmoid, \ 7 | Softplus, Softshrink, MultiheadAttention, PReLU, Softsign, Softmin, Tanhshrink, RReLU, GLU 8 | from .loss import L1Loss, NLLLoss, KLDivLoss, MSELoss, BCELoss, BCEWithLogitsLoss, NLLLoss2d, \ 9 | CosineEmbeddingLoss, CTCLoss, HingeEmbeddingLoss, MarginRankingLoss, \ 10 | MultiLabelMarginLoss, MultiLabelSoftMarginLoss, MultiMarginLoss, \ 11 | SmoothL1Loss, SoftMarginLoss, CrossEntropyLoss, TripletMarginLoss, PoissonNLLLoss 12 | from .container import Container, Sequential, ModuleList, ModuleDict, ParameterList, ParameterDict 13 | from .pooling import AvgPool1d, AvgPool2d, AvgPool3d, MaxPool1d, MaxPool2d, MaxPool3d, \ 14 | MaxUnpool1d, MaxUnpool2d, MaxUnpool3d, FractionalMaxPool2d, FractionalMaxPool3d, LPPool1d, LPPool2d, \ 15 | AdaptiveMaxPool1d, AdaptiveMaxPool2d, AdaptiveMaxPool3d, AdaptiveAvgPool1d, AdaptiveAvgPool2d, AdaptiveAvgPool3d 16 | from .batchnorm import BatchNorm1d, BatchNorm2d, BatchNorm3d, SyncBatchNorm 17 | from .instancenorm import InstanceNorm1d, InstanceNorm2d, InstanceNorm3d 18 | from .normalization import LocalResponseNorm, CrossMapLRN2d, LayerNorm, GroupNorm 19 | from .dropout import Dropout, Dropout2d, Dropout3d, AlphaDropout, FeatureAlphaDropout 20 | from .padding import ReflectionPad1d, ReflectionPad2d, ReplicationPad1d, ReplicationPad2d, \ 21 | ReplicationPad3d, ZeroPad2d, ConstantPad1d, ConstantPad2d, ConstantPad3d 22 | from .sparse import Embedding, EmbeddingBag 23 | from .rnn import RNNBase, RNN, LSTM, GRU, \ 24 | RNNCellBase, RNNCell, LSTMCell, GRUCell 25 | from .pixelshuffle import PixelShuffle 26 | from .upsampling import UpsamplingNearest2d, UpsamplingBilinear2d, Upsample 27 | from .distance import PairwiseDistance, CosineSimilarity 28 | from .fold import Fold, Unfold 29 | from .adaptive import AdaptiveLogSoftmaxWithLoss 30 | from .transformer import TransformerEncoder, TransformerDecoder, \ 31 | TransformerEncoderLayer, TransformerDecoderLayer, Transformer 32 | from .flatten import Flatten 33 | from .lala import Wrapper 34 | 35 | __all__ = [ 36 | 'Module', 'Identity', 'Linear', 'Conv1d','Conv2d', 'Conv3d', 'ConvTranspose1d', 37 | 'ConvTranspose2d', 'ConvTranspose3d', 'Threshold', 'ReLU', 'Hardtanh', 'ReLU6', 38 | 'Sigmoid', 'Tanh', 'Softmax', 'Softmax2d', 'LogSoftmax', 'ELU', 'SELU', 'CELU', 'GLU', 'GELU', 'Hardshrink', 39 | 'LeakyReLU', 'LogSigmoid', 'Softplus', 'Softshrink', 'MultiheadAttention', 'PReLU', 'Softsign', 'Softmin', 40 | 'Tanhshrink', 'RReLU', 'L1Loss', 'NLLLoss', 'KLDivLoss', 'MSELoss', 'BCELoss', 'BCEWithLogitsLoss', 41 | 'NLLLoss2d', 'PoissonNLLLoss', 'CosineEmbeddingLoss', 'CTCLoss', 'HingeEmbeddingLoss', 'MarginRankingLoss', 42 | 'MultiLabelMarginLoss', 'MultiLabelSoftMarginLoss', 'MultiMarginLoss', 'SmoothL1Loss', 43 | 'SoftMarginLoss', 'CrossEntropyLoss', 'Container', 'Sequential', 'ModuleList', 'ModuleDict', 44 | 'ParameterList', 'ParameterDict', 'AvgPool1d', 'AvgPool2d', 'AvgPool3d', 'MaxPool1d', 'MaxPool2d', 45 | 'MaxPool3d', 'MaxUnpool1d', 'MaxUnpool2d', 'MaxUnpool3d', 'FractionalMaxPool2d', "FractionalMaxPool3d", 46 | 'LPPool1d', 'LPPool2d', 'LocalResponseNorm', 'BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d', 'InstanceNorm1d', 47 | 'InstanceNorm2d', 'InstanceNorm3d', 'LayerNorm', 'GroupNorm', 'SyncBatchNorm', 48 | 'Dropout', 'Dropout2d', 'Dropout3d', 'AlphaDropout', 'FeatureAlphaDropout', 49 | 'ReflectionPad1d', 'ReflectionPad2d', 'ReplicationPad2d', 'ReplicationPad1d', 'ReplicationPad3d', 50 | 'CrossMapLRN2d', 'Embedding', 'EmbeddingBag', 'RNNBase', 'RNN', 'LSTM', 'GRU', 'RNNCellBase', 'RNNCell', 51 | 'LSTMCell', 'GRUCell', 'PixelShuffle', 'Upsample', 'UpsamplingNearest2d', 'UpsamplingBilinear2d', 52 | 'PairwiseDistance', 'AdaptiveMaxPool1d', 'AdaptiveMaxPool2d', 'AdaptiveMaxPool3d', 'AdaptiveAvgPool1d', 53 | 'AdaptiveAvgPool2d', 'AdaptiveAvgPool3d', 'TripletMarginLoss', 'ZeroPad2d', 'ConstantPad1d', 'ConstantPad2d', 54 | 'ConstantPad3d', 'Bilinear', 'CosineSimilarity', 'Unfold', 'Fold', 55 | 'AdaptiveLogSoftmaxWithLoss', 'TransformerEncoder', 'TransformerDecoder', 56 | 'TransformerEncoderLayer', 'TransformerDecoderLayer', 'Transformer', 57 | 'Flatten', 'Wrapper' 58 | ] 59 | -------------------------------------------------------------------------------- /nn/utils/weight_norm.py: -------------------------------------------------------------------------------- 1 | r""" 2 | Weight Normalization from https://arxiv.org/abs/1602.07868 3 | """ 4 | from torch.nn.parameter import Parameter 5 | from torch import _weight_norm, norm_except_dim 6 | 7 | 8 | class WeightNorm(object): 9 | def __init__(self, name, dim): 10 | if dim is None: 11 | dim = -1 12 | self.name = name 13 | self.dim = dim 14 | 15 | def compute_weight(self, module): 16 | g = getattr(module, self.name + '_g') 17 | v = getattr(module, self.name + '_v') 18 | return _weight_norm(v, g, self.dim) 19 | 20 | @staticmethod 21 | def apply(module, name, dim): 22 | for k, hook in module._forward_pre_hooks.items(): 23 | if isinstance(hook, WeightNorm) and hook.name == name: 24 | raise RuntimeError("Cannot register two weight_norm hooks on " 25 | "the same parameter {}".format(name)) 26 | 27 | if dim is None: 28 | dim = -1 29 | 30 | fn = WeightNorm(name, dim) 31 | 32 | weight = getattr(module, name) 33 | 34 | # remove w from parameter list 35 | del module._parameters[name] 36 | 37 | # add g and v as new parameters and express w as g/||v|| * v 38 | module.register_parameter(name + '_g', Parameter(norm_except_dim(weight, 2, dim).data)) 39 | module.register_parameter(name + '_v', Parameter(weight.data)) 40 | setattr(module, name, fn.compute_weight(module)) 41 | 42 | # recompute weight before every forward() 43 | module.register_forward_pre_hook(fn) 44 | 45 | return fn 46 | 47 | def remove(self, module): 48 | weight = self.compute_weight(module) 49 | delattr(module, self.name) 50 | del module._parameters[self.name + '_g'] 51 | del module._parameters[self.name + '_v'] 52 | module.register_parameter(self.name, Parameter(weight.data)) 53 | 54 | def __call__(self, module, inputs): 55 | setattr(module, self.name, self.compute_weight(module)) 56 | 57 | 58 | def weight_norm(module, name='weight', dim=0): 59 | r"""Applies weight normalization to a parameter in the given module. 60 | 61 | .. math:: 62 | \mathbf{w} = g \dfrac{\mathbf{v}}{\|\mathbf{v}\|} 63 | 64 | Weight normalization is a reparameterization that decouples the magnitude 65 | of a weight tensor from its direction. This replaces the parameter specified 66 | by :attr:`name` (e.g. ``'weight'``) with two parameters: one specifying the magnitude 67 | (e.g. ``'weight_g'``) and one specifying the direction (e.g. ``'weight_v'``). 68 | Weight normalization is implemented via a hook that recomputes the weight 69 | tensor from the magnitude and direction before every :meth:`~Module.forward` 70 | call. 71 | 72 | By default, with ``dim=0``, the norm is computed independently per output 73 | channel/plane. To compute a norm over the entire weight tensor, use 74 | ``dim=None``. 75 | 76 | See https://arxiv.org/abs/1602.07868 77 | 78 | Args: 79 | module (Module): containing module 80 | name (str, optional): name of weight parameter 81 | dim (int, optional): dimension over which to compute the norm 82 | 83 | Returns: 84 | The original module with the weight norm hook 85 | 86 | Example:: 87 | 88 | >>> m = weight_norm(nn.Linear(20, 40), name='weight') 89 | >>> m 90 | Linear(in_features=20, out_features=40, bias=True) 91 | >>> m.weight_g.size() 92 | torch.Size([40, 1]) 93 | >>> m.weight_v.size() 94 | torch.Size([40, 20]) 95 | 96 | """ 97 | WeightNorm.apply(module, name, dim) 98 | return module 99 | 100 | 101 | def remove_weight_norm(module, name='weight'): 102 | r"""Removes the weight normalization reparameterization from a module. 103 | 104 | Args: 105 | module (Module): containing module 106 | name (str, optional): name of weight parameter 107 | 108 | Example: 109 | >>> m = weight_norm(nn.Linear(20, 40)) 110 | >>> remove_weight_norm(m) 111 | """ 112 | for k, hook in module._forward_pre_hooks.items(): 113 | if isinstance(hook, WeightNorm) and hook.name == name: 114 | hook.remove(module) 115 | del module._forward_pre_hooks[k] 116 | return module 117 | 118 | raise ValueError("weight_norm of '{}' not found in {}" 119 | .format(name, module)) 120 | -------------------------------------------------------------------------------- /nn/modules/module.pyi: -------------------------------------------------------------------------------- 1 | from ... import Tensor, device, dtype 2 | from .. import Parameter 3 | from typing import Union, Tuple, Any, Callable, Iterator, Set, Optional, overload, TypeVar, Mapping, Dict, Generic 4 | from collections import OrderedDict 5 | from ...utils.hooks import RemovableHandle 6 | 7 | _grad_t = Union[Tuple[Tensor, ...], Tensor] 8 | # See https://mypy.readthedocs.io/en/latest/generics.html#generic-methods-and-generic-self for the use 9 | # of `T` to annotate `self`. Many methods of `Module` return `self` and we want those return values to be 10 | # the type of the subclass, not the looser type of `Module`. 11 | T = TypeVar('T') 12 | # We parameter modules by the return type of its `forward` (and therefore `__call__`) method. This allows 13 | # type inference to infer that the return value of calling a module in the canonical way (via `__call__)` is the 14 | # same as the custom `forward` function of the submodule. Submodules tha wish to opt in this functionality be 15 | # defined as eg class ReturnsTwoTensors(Module[Tuple[Tensor, Tensor]]): ... 16 | T_co = TypeVar('T_co', covariant=True) 17 | 18 | 19 | class Module(Generic[T_co]): 20 | def __init__(self) -> None: ... 21 | 22 | def forward(self, *input: Any, **kwargs: Any) -> T_co: ... # type: ignore 23 | 24 | def __call__(self, *input: Any, **kwargs: Any) -> T_co: ... # type: ignore 25 | 26 | def register_buffer(self, name: str, tensor: Tensor) -> None: ... 27 | 28 | def register_parameter(self, name: str, param: Parameter) -> None: ... 29 | 30 | def add_module(self, name: str, module: 'Module') -> None: ... 31 | 32 | def apply(self: T, fn: Callable[['Module'], None]) -> T: ... 33 | 34 | def cuda(self: T, device: Optional[Union[int, device]] = ...) -> T: ... 35 | 36 | def cpu(self: T) -> T: ... 37 | 38 | def type(self: T, dst_type: Union[dtype, str]) -> T: ... 39 | 40 | def float(self: T) -> T: ... 41 | 42 | def double(self: T) -> T: ... 43 | 44 | def half(self: T) -> T: ... 45 | 46 | @overload 47 | def to(self: T, device: Optional[Union[int, device]] = ..., dtype: Optional[Union[dtype, str]] = ..., 48 | non_blocking: bool = ...) -> T: ... 49 | 50 | @overload 51 | def to(self: T, dtype: Union[dtype, str], non_blocking: bool = ...) -> T: ... 52 | 53 | @overload 54 | def to(self: T, tensor: Tensor, non_blocking: bool = ...) -> T: ... 55 | 56 | def register_backward_hook(self, hook: Callable[ 57 | ['Module', _grad_t, _grad_t], Union[None, Tensor]]) -> RemovableHandle: ... 58 | 59 | # The hook takes a module as a first argument and variadic arguments after that, but there is no way to express that 60 | def register_forward_pre_hook(self, hook: Callable[..., None]) -> RemovableHandle: ... 61 | 62 | def register_forward_hook(self, hook: Callable[..., None]) -> RemovableHandle: ... 63 | 64 | def __getattr__(self, name: str) -> Union[Tensor, 'Module']: ... 65 | 66 | # TODO double-check this 67 | def __setattr__(self, name: str, value: Union[Tensor, 'Module']) -> None: ... 68 | 69 | # The user can pass an optional arbitrary mappable object to `state_dict`, in which case `state_dict` returns 70 | # back that same object. But if they pass nothing, an `OrederedDict` is created and returned. 71 | T_destination = TypeVar('T_destination', bound=Mapping[str, Tensor]) 72 | 73 | @overload 74 | def state_dict(self, destination: T_destination, prefix: str = ..., keep_vars: bool = ...) -> T_destination: ... 75 | 76 | @overload 77 | def state_dict(self, prefix: str = ..., keep_vars: bool = ...) -> OrderedDict[str, Tensor]: ... 78 | 79 | def load_state_dict(self, state_dict: Union[Dict[str, Tensor], OrderedDict[str, Tensor]], strict: bool = ...): ... 80 | 81 | def parameters(self, recurse: bool = ...) -> Iterator[Parameter]: ... 82 | 83 | def named_parameters(self, prefix: str = ..., recurse: bool = ...) -> Iterator[Tuple[str, Parameter]]: ... 84 | 85 | def buffers(self, recurse: bool = ...) -> Iterator[Tensor]: ... 86 | 87 | def named_buffers(self, prefix: str = ..., recurse: bool = ...) -> Iterator[Tuple[str, Tensor]]: ... 88 | 89 | def children(self) -> Iterator['Module']: ... 90 | 91 | def named_children(self) -> Iterator[Tuple[str, 'Module']]: ... 92 | 93 | def modules(self) -> Iterator['Module']: ... 94 | 95 | def named_modules(self, memo: Optional[Set['Module']] = ..., prefix: str = ...) -> Iterator[ 96 | Tuple[str, 'Module']]: ... 97 | 98 | def train(self: T, mode: bool = ...) -> T: ... 99 | 100 | def eval(self: T) -> T: ... 101 | 102 | def zero_grad(self) -> None: ... 103 | 104 | def share_memory(self: T) -> T: ... 105 | 106 | def extra_repr(self) -> str: ... 107 | -------------------------------------------------------------------------------- /nn/modules/conv.pyi: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from typing import Any, Optional, List, Tuple, Union 3 | from ... import Tensor 4 | from ..common_types import _size_1_t, _size_2_t, _size_3_t 5 | 6 | 7 | class _ConvNd(Module): 8 | in_channels: int = ... 9 | out_channels: int = ... 10 | kernel_size: Tuple[int, ...] = ... 11 | stride: Tuple[int, ...] = ... 12 | padding: Tuple[int, ...] = ... 13 | dilation: Tuple[int, ...] = ... 14 | transposed: bool = ... 15 | output_padding: Tuple[int, ...] = ... 16 | groups: int = ... 17 | padding_mode: str = ... 18 | weight: Tensor = ... 19 | bias: Tensor = ... 20 | 21 | # padding_mode can only one of an enumerated set of strings. Python typing will eventually support precisely typing 22 | # this with the `Literal` type. 23 | def __init__(self, in_channels: Any, out_channels: Any, kernel_size: Any, stride: Any, padding: Any, dilation: Any, 24 | transposed: Any, output_padding: Any, groups: Any, bias: Any, padding_mode: Any) -> None: ... 25 | 26 | def reset_parameters(self) -> None: ... 27 | 28 | 29 | class Conv1d(_ConvNd): 30 | def __init__(self, in_channels: int, out_channels: int, kernel_size: _size_1_t, stride: _size_1_t = ..., 31 | padding: _size_1_t = ..., dilation: _size_1_t = ..., groups: int = ..., bias: bool = ..., 32 | padding_mode: str = ...) -> None: ... 33 | 34 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 35 | 36 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 37 | 38 | 39 | class Conv2d(_ConvNd): 40 | def __init__(self, in_channels: int, out_channels: int, kernel_size: _size_2_t, stride: _size_2_t = ..., 41 | padding: _size_2_t = ..., dilation: _size_2_t = ..., groups: int = ..., bias: bool = ..., 42 | padding_mode: str = ...) -> None: ... 43 | 44 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 45 | 46 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 47 | 48 | 49 | class Conv3d(_ConvNd): 50 | def __init__(self, in_channels: int, out_channels: int, kernel_size: _size_3_t, stride: _size_3_t = ..., 51 | padding: _size_3_t = ..., dilation: _size_3_t = ..., groups: int = ..., bias: bool = ..., 52 | padding_mode: str = ...) -> None: ... 53 | 54 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 55 | 56 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 57 | 58 | 59 | class _ConvTransposeMixin: 60 | def forward(self, input: Tensor, output_size: Optional[List[int]] = ...): ... # type: ignore 61 | def __call__(self, input: Tensor, output_size: Optional[List[int]] = ...): ... # type: ignore 62 | 63 | # We need a '# type: ignore' at the end of the declaration of each class that inherits from 64 | # `_ConvTransposeMixin` since the `forward` method declared in `_ConvTransposeMixin` is 65 | # incompatible with the `forward` method declared in `Module`. 66 | class ConvTranspose1d(_ConvTransposeMixin, _ConvNd): # type: ignore 67 | def __init__(self, in_channels: int, out_channels: int, kernel_size: _size_1_t, stride: _size_1_t = ..., 68 | padding: _size_1_t = ..., output_padding: _size_1_t = ..., groups: int = ..., bias: bool = ..., 69 | dilation: int = ..., padding_mode: str = ...) -> None: ... 70 | 71 | def forward(self, input: Tensor, output_size: Optional[List[int]] = ...) -> Tensor: ... # type: ignore 72 | 73 | def __call__(self, input: Tensor, output_size: Optional[List[int]] = ...) -> Tensor: ... # type: ignore 74 | 75 | 76 | class ConvTranspose2d(_ConvTransposeMixin, _ConvNd): # type: ignore 77 | def __init__(self, in_channels: int, out_channels: int, kernel_size: _size_2_t, stride: _size_2_t = ..., 78 | padding: _size_2_t = ..., output_padding: _size_2_t = ..., groups: int = ..., bias: bool = ..., 79 | dilation: int = ..., padding_mode: str = ...) -> None: ... 80 | 81 | def forward(self, input: Tensor, output_size: Optional[List[int]] = ...) -> Tensor: ... # type: ignore 82 | 83 | def __call__(self, input: Tensor, output_size: Optional[List[int]] = ...) -> Tensor: ... # type: ignore 84 | 85 | 86 | class ConvTranspose3d(_ConvTransposeMixin, _ConvNd): # type: ignore 87 | def __init__(self, in_channels: int, out_channels: int, kernel_size: _size_3_t, stride: _size_3_t = ..., 88 | padding: _size_3_t = ..., output_padding: _size_3_t = ..., groups: int = ..., bias: bool = ..., 89 | dilation: int = ..., padding_mode: str = ...) -> None: ... 90 | 91 | def forward(self, input: Tensor, output_size: Optional[List[int]] = ...) -> Tensor: ... # type: ignore 92 | 93 | def __call__(self, input: Tensor, output_size: Optional[List[int]] = ...) -> Tensor: ... # type: ignore 94 | -------------------------------------------------------------------------------- /nn/parallel/_functions.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | 3 | import torch 4 | import torch.cuda.comm as comm 5 | from torch.autograd import Function 6 | from torch.cuda._utils import _get_device_index 7 | 8 | 9 | class Broadcast(Function): 10 | 11 | @staticmethod 12 | def forward(ctx, target_gpus, *inputs): 13 | if not all(input.is_cuda for input in inputs): 14 | raise TypeError('Broadcast function not implemented for CPU tensors') 15 | target_gpus = list(map(lambda x: _get_device_index(x, True), target_gpus)) 16 | ctx.target_gpus = target_gpus 17 | if len(inputs) == 0: 18 | return tuple() 19 | ctx.num_inputs = len(inputs) 20 | ctx.input_device = inputs[0].get_device() 21 | outputs = comm.broadcast_coalesced(inputs, ctx.target_gpus) 22 | non_differentiables = [] 23 | for idx, input_requires_grad in enumerate(ctx.needs_input_grad[1:]): 24 | if not input_requires_grad: 25 | for output in outputs: 26 | non_differentiables.append(output[idx]) 27 | ctx.mark_non_differentiable(*non_differentiables) 28 | return tuple([t for tensors in outputs for t in tensors]) 29 | 30 | @staticmethod 31 | def backward(ctx, *grad_outputs): 32 | return (None,) + ReduceAddCoalesced.apply(ctx.input_device, ctx.num_inputs, *grad_outputs) 33 | 34 | 35 | class ReduceAddCoalesced(Function): 36 | 37 | @staticmethod 38 | def forward(ctx, destination, num_inputs, *grads): 39 | ctx.target_gpus = [grads[i].get_device() for i in range(0, len(grads), num_inputs)] 40 | 41 | grads = [grads[i:i + num_inputs] 42 | for i in range(0, len(grads), num_inputs)] 43 | return comm.reduce_add_coalesced(grads, destination) 44 | 45 | @staticmethod 46 | def backward(ctx, *grad_outputs): 47 | return (None, None,) + Broadcast.apply(ctx.target_gpus, *grad_outputs) 48 | 49 | 50 | class Gather(Function): 51 | 52 | @staticmethod 53 | def forward(ctx, target_device, dim, *inputs): 54 | assert all(map(lambda i: i.is_cuda, inputs)) 55 | target_device = _get_device_index(target_device, True) 56 | ctx.target_device = target_device 57 | ctx.dim = dim 58 | ctx.input_gpus = tuple(map(lambda i: i.get_device(), inputs)) 59 | if all(t.dim() == 0 for t in inputs) and dim == 0: 60 | inputs = tuple(t.view(1) for t in inputs) 61 | warnings.warn('Was asked to gather along dimension 0, but all ' 62 | 'input tensors were scalars; will instead unsqueeze ' 63 | 'and return a vector.') 64 | ctx.unsqueezed_scalar = True 65 | else: 66 | ctx.unsqueezed_scalar = False 67 | ctx.input_sizes = tuple(map(lambda i: i.size(ctx.dim), inputs)) 68 | return comm.gather(inputs, ctx.dim, ctx.target_device) 69 | 70 | @staticmethod 71 | def backward(ctx, grad_output): 72 | scattered_grads = Scatter.apply(ctx.input_gpus, ctx.input_sizes, ctx.dim, grad_output) 73 | if ctx.unsqueezed_scalar: 74 | scattered_grads = tuple(g[0] for g in scattered_grads) 75 | return (None, None) + scattered_grads 76 | 77 | 78 | class Scatter(Function): 79 | 80 | @staticmethod 81 | def forward(ctx, target_gpus, chunk_sizes, dim, input): 82 | target_gpus = list(map(lambda x: _get_device_index(x, True), target_gpus)) 83 | ctx.dim = dim 84 | ctx.input_device = input.get_device() if input.is_cuda else -1 85 | streams = None 86 | if ctx.input_device == -1: 87 | # Perform CPU to GPU copies in a background stream 88 | streams = [_get_stream(device) for device in target_gpus] 89 | outputs = comm.scatter(input, target_gpus, chunk_sizes, ctx.dim, streams) 90 | # Synchronize with the copy stream 91 | if streams is not None: 92 | for i, output in enumerate(outputs): 93 | with torch.cuda.device(target_gpus[i]): 94 | main_stream = torch.cuda.current_stream() 95 | main_stream.wait_stream(streams[i]) 96 | output.record_stream(main_stream) 97 | return outputs 98 | 99 | @staticmethod 100 | def backward(ctx, *grad_output): 101 | return None, None, None, Gather.apply(ctx.input_device, ctx.dim, *grad_output) 102 | 103 | 104 | # background streams used for copying 105 | _streams = None 106 | 107 | 108 | def _get_stream(device): 109 | """Gets a background stream for copying between CPU and GPU""" 110 | global _streams 111 | if device == -1: 112 | return None 113 | if _streams is None: 114 | _streams = [None] * torch.cuda.device_count() 115 | if _streams[device] is None: 116 | _streams[device] = torch.cuda.Stream(device) 117 | return _streams[device] 118 | -------------------------------------------------------------------------------- /quantization/_quantize_script.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | 3 | import torch 4 | from .qconfig import QConfig 5 | 6 | class ConvPackedParams(torch.nn.Module): 7 | def __init__(self): 8 | super(ConvPackedParams, self).__init__() 9 | wq = torch._empty_affine_quantized([1, 1, 1, 1], scale=1.0, zero_point=0, dtype=torch.qint8) 10 | self.stride = [1, 1] 11 | self.padding = [0, 0] 12 | self.dilation = [1, 1] 13 | self.groups = 1 14 | self.set_weight_bias(wq, None) 15 | 16 | @torch.jit.export 17 | def set_conv_params(self, stride, padding, dilation, groups): 18 | # type: (List[int], List[int], List[int], int) -> None 19 | self.stride = stride 20 | self.padding = padding 21 | self.dilation = dilation 22 | self.groups = groups 23 | 24 | @torch.jit.export 25 | def set_weight_bias(self, weight, bias): 26 | # type: (torch.Tensor, Optional[torch.Tensor]) -> None 27 | self._packed_params = torch.ops.quantized.conv2d_prepack(weight, bias, self.stride, self.padding, self.dilation, self.groups) 28 | 29 | @torch.jit.export 30 | def _weight_bias(self): 31 | return torch.ops.quantized.conv2d_unpack(self._packed_params) 32 | 33 | def forward(self, x): 34 | return x 35 | 36 | @torch.jit.export 37 | def __getstate__(self): 38 | qweight, bias = self._weight_bias() 39 | return (qweight, 40 | bias, 41 | self.stride, 42 | self.padding, 43 | self.dilation, 44 | self.groups, 45 | self.training) 46 | 47 | @torch.jit.export 48 | def __setstate__(self, state): 49 | self.stride = state[2] 50 | self.padding = state[3] 51 | self.dilation = state[4] 52 | self.groups = state[5] 53 | self.set_weight_bias(state[0], 54 | state[1]) 55 | self.training = state[6] 56 | 57 | linear_packed_params = None 58 | conv_packed_params = None 59 | if 'fbgemm' in torch.backends.quantized.supported_engines: 60 | linear_packed_params = torch.jit.script(torch.nn.quantized.modules.linear.LinearPackedParams())._c 61 | conv_packed_params = torch.jit.script(ConvPackedParams())._c 62 | 63 | def _check_is_script_module(model): 64 | if not isinstance(model, torch.jit.ScriptModule): 65 | raise ValueError('input must be a script module, got: ' + str(type(model))) 66 | 67 | def prepare_script(model, qconfig_dict, inplace=False): 68 | _check_is_script_module(model) 69 | if not inplace: 70 | model = model.copy() 71 | torch._C._jit_pass_insert_observers(model._c, 72 | 'forward', 73 | qconfig_dict, 74 | True) 75 | return model 76 | 77 | def convert_script(model, inplace=False): 78 | _check_is_script_module(model) 79 | if not inplace: 80 | model = model.copy() 81 | torch._C._jit_pass_insert_quant_dequant(model._c, 'forward', True) 82 | if 'fbgemm' in torch.backends.quantized.supported_engines: 83 | torch._C._jit_pass_insert_prepack_unpack(model._c) 84 | if linear_packed_params and conv_packed_params: 85 | torch._C._jit_pass_fold_prepack(model._c, 86 | linear_packed_params, 87 | conv_packed_params) 88 | 89 | return model 90 | 91 | # TODO: non-scriptable QConfig will be supported later 92 | def script_qconfig(qconfig): 93 | return QConfig( 94 | activation=torch.jit.script(qconfig.activation())._c, 95 | weight=torch.jit.script(qconfig.weight())._c) 96 | 97 | def quantize_script(model, qconfig_dict, run_fn, run_args, inplace=False): 98 | _check_is_script_module(model) 99 | if not model._c._has_method('forward'): 100 | raise ValueError('input script module does not have forward method') 101 | assert not inplace, "We don't support inplace right now" 102 | if not inplace: 103 | model = model.copy() 104 | scripted_qconfig_dict = {k: script_qconfig(v) for k, v in qconfig_dict.items()} 105 | # We are not going to run fold_convbn pass right now 106 | # since it is not able to work correctly, we will 107 | # revisit after constants is properly handled in 108 | # JIT 109 | # torch._C._jit_pass_fold_convbn(model._c) 110 | prepare_script(model, scripted_qconfig_dict, True) 111 | run_fn(model._c._get_method('forward'), *run_args) 112 | # When we mutating graph we didn't create a new ClassType 113 | # and the graph executor will run an out dated version 114 | # of the graph if we do inplace graph mutation, therefore 115 | # we copy the model here 116 | # [TODO] This will be fixed later when we figure out 117 | # how to properly mutate types 118 | model = convert_script(model, False) 119 | return model 120 | -------------------------------------------------------------------------------- /quantization/qconfig.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | from collections import namedtuple 3 | from .observer import * 4 | from .fake_quantize import * 5 | import torch.nn as nn 6 | 7 | class QConfig(namedtuple('QConfig', ['activation', 'weight'])): 8 | """ 9 | Describes how to quantize a layer or a part of the network by providing 10 | settings (observer classes) for activations and weights respectively. 11 | 12 | 13 | Note that QConfig needs to contain observer **classes** (like MinMaxObserver) or a callable that returns 14 | instances on invocation, not the concrete observer instances themselves. 15 | Quantization preparation function will instantiate observers multiple times for each of the layers. 16 | 17 | 18 | Observer classes have usually reasonable default arguments, but they can be overwritten with `with_args` 19 | method (that behaves like functools.partial): 20 | 21 | my_qconfig = QConfig(activation=MinMaxObserver.with_args(dtype=torch.qint8), 22 | weight=default_observer.with_args(dtype=torch.qint8)) 23 | """ 24 | def __new__(cls, activation, weight): 25 | # catch common mistakes 26 | if isinstance(activation, nn.Module) or isinstance(weight, nn.Module): 27 | raise ValueError("QConfig received observer instance, please pass observer class instead. " + 28 | "Use MyObserver.with_args(x=1) to override arguments to constructor if needed") 29 | return super(QConfig, cls).__new__(cls, activation, weight) 30 | 31 | 32 | default_qconfig = QConfig(activation=default_observer, 33 | weight=default_weight_observer) 34 | 35 | default_debug_qconfig = QConfig(weight=default_weight_observer, 36 | activation=default_debug_observer) 37 | 38 | default_per_channel_qconfig = QConfig(activation=default_observer, 39 | weight=default_per_channel_weight_observer) 40 | 41 | class QConfigDynamic(namedtuple('QConfigDynamic', ['weight'])): 42 | """ 43 | Describes how to dynamically quantize a layer or a part of the network by providing 44 | settings (observer classe) for weights. 45 | 46 | It's like QConfig, but for dynamic quantization. 47 | 48 | Note that QConfigDynamic needs to contain observer **classes** (like MinMaxObserver) or a callable that returns 49 | instances on invocation, not the concrete observer instances themselves. 50 | Quantization function will instantiate observers multiple times for each of the layers. 51 | 52 | Observer classes have usually reasonable default arguments, but they can be overwritten with `with_args` 53 | method (that behaves like functools.partial): 54 | 55 | my_qconfig = QConfigDynamic(weight=default_observer.with_args(dtype=torch.qint8)) 56 | """ 57 | def __new__(cls, weight): 58 | # catch common mistakes 59 | if isinstance(weight, nn.Module): 60 | raise ValueError("QConfigDynamic received observer instance, please pass observer class instead. " + 61 | "Use MyObserver.with_args(x=1) to override arguments to constructor if needed") 62 | return super(QConfigDynamic, cls).__new__(cls, weight) 63 | 64 | default_dynamic_qconfig = QConfigDynamic(weight=default_weight_observer) 65 | float16_dynamic_qconfig = QConfigDynamic(weight=NoopObserver.with_args(dtype=torch.float16)) 66 | per_channel_dynamic_qconfig = QConfigDynamic(weight=default_per_channel_weight_observer) 67 | 68 | default_qat_qconfig = QConfig(activation=default_fake_quant, 69 | weight=default_weight_fake_quant) 70 | 71 | default_weight_only_qconfig = QConfig(activation=torch.nn.Identity, 72 | weight=default_weight_fake_quant) 73 | default_activation_only_qconfig = QConfig(activation=default_fake_quant, 74 | weight=torch.nn.Identity) 75 | 76 | def get_default_qconfig(backend='fbgemm'): 77 | if backend == 'fbgemm': 78 | qconfig = QConfig(activation=HistogramObserver.with_args(reduce_range=True), 79 | weight=default_per_channel_weight_observer) 80 | elif backend == 'qnnpack': 81 | qconfig = QConfig(activation=HistogramObserver.with_args(reduce_range=False), 82 | weight=default_weight_observer) 83 | else: 84 | raise ValueError("Unknown backend, please specify qconfig manually") 85 | return qconfig 86 | 87 | def get_default_qat_qconfig(backend='fbgemm'): 88 | # Histogram observer is too slow for quantization aware training 89 | if backend == 'fbgemm': 90 | qconfig = QConfig(activation=FakeQuantize.with_args(observer=MovingAverageMinMaxObserver, 91 | quant_min=0, 92 | quant_max=255, 93 | reduce_range=True), 94 | weight=default_per_channel_weight_fake_quant) 95 | elif backend == 'qnnpack': 96 | qconfig = QConfig(activation=FakeQuantize.with_args(observer=MovingAverageMinMaxObserver, 97 | quant_min=0, 98 | quant_max=255, 99 | reduce_range=False), 100 | weight=default_weight_fake_quant) 101 | else: 102 | raise ValueError("Unknown backend, please specify qconfig manually") 103 | 104 | return qconfig 105 | -------------------------------------------------------------------------------- /nn/modules/rnn.pyi: -------------------------------------------------------------------------------- 1 | from ..parameter import Parameter 2 | from .module import Module 3 | from typing import Any, Optional, Tuple, List 4 | from ... import Tensor 5 | 6 | 7 | def apply_permutation(tensor: Tensor, permutation: Tensor, dim: int = ...) -> Tensor: ... 8 | 9 | 10 | class RNNBase(Module): 11 | mode: str = ... 12 | input_size: int = ... 13 | hidden_size: int = ... 14 | num_layers: int = ... 15 | bias: bool = ... 16 | batch_first: bool = ... 17 | dropout: float = ... 18 | bidirectional: bool = ... 19 | 20 | def __init__(self, mode: str, input_size: int, hidden_size: int, num_layers: int = ..., bias: bool = ..., 21 | batch_first: bool = ..., dropout: float = ..., bidirectional: bool = ...) -> None: ... 22 | 23 | def flatten_parameters(self) -> List[Parameter]: ... 24 | 25 | def reset_parameters(self) -> None: ... 26 | 27 | def get_flat_weights(self): ... 28 | 29 | def check_input(self, input: Tensor, batch_sizes: Optional[Tensor]) -> None: ... 30 | 31 | def get_expected_hidden_size(self, input: Tensor, batch_sizes: Optional[Tensor]) -> Tuple[int, int, int]: ... 32 | 33 | def check_hidden_size(self, hx: Tensor, expected_hidden_size: Tuple[int, int, int], msg: str = ...) -> None: ... 34 | 35 | def check_forward_args(self, input: Any, hidden: Any, batch_sizes: Optional[Tensor]) -> None: ... 36 | 37 | def permute_hidden(self, hx: Any, permutation: Any): ... 38 | 39 | def forward(self, input: Tensor, hx: Optional[Any] = ...) -> Any: ... # type: ignore 40 | 41 | def __call__(self, input: Tensor, hx: Optional[Any] = ...) -> Any: ... # type: ignore 42 | 43 | @property 44 | def all_weights(self) -> List[Parameter]: ... 45 | 46 | 47 | class RNN(RNNBase): 48 | def __init__(self, input_size: int, hidden_size: int, num_layers: int = ..., bias: bool = ..., 49 | batch_first: bool = ..., dropout: float = ..., bidirectional: bool = ..., 50 | nonlinearity: str = ...) -> None: ... 51 | 52 | def forward(self, input: Tensor, hx: Optional[Tensor] = ...) -> Tensor: ... # type: ignore 53 | 54 | def __call__(self, input: Tensor, hx: Optional[Tensor] = ...) -> Tensor: ... # type: ignore 55 | 56 | 57 | class LSTM(RNNBase): 58 | def __init__(self, input_size: int, hidden_size: int, num_layers: int = ..., bias: bool = ..., 59 | batch_first: bool = ..., dropout: float = ..., bidirectional: bool = ..., 60 | nonlinearity: str = ...) -> None: ... 61 | 62 | def check_forward_args(self, input: Tensor, hidden: Tuple[Tensor, Tensor], 63 | batch_sizes: Optional[Tensor]) -> None: ... 64 | 65 | def permute_hidden(self, hx: Tuple[Tensor, Tensor], permutation: Optional[Tensor]) -> Tuple[Tensor, Tensor]: ... 66 | 67 | def forward_impl(self, input: Tensor, hx: Optional[Tuple[Tensor, Tensor]], batch_sizes: Optional[Tensor], 68 | max_batch_size: int, sorted_indices: Optional[Tensor]) -> Tuple[Tensor, Tuple[Tensor, Tensor]]: ... 69 | 70 | def forward_tensor(self, input: Tensor, hx: Optional[Tuple[Tensor, Tensor]] = ...) -> Tuple[ 71 | Tensor, Tuple[Tensor, Tensor]]: ... 72 | 73 | def forward_packed(self, input: Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]], 74 | hx: Optional[Tuple[Tensor, Tensor]] = ...) -> Tuple[ 75 | Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]], Tuple[Tensor, Tensor]]: ... 76 | 77 | def forward(self, input: Tensor, hx: Optional[Tuple[Tensor, Tensor]] = ...) -> Tuple[Tensor, Tuple[Tensor, Tensor]]: ... # type: ignore 78 | 79 | def __call__(self, input: Tensor, hx: Optional[Tuple[Tensor, Tensor]] = ...) -> Tuple[Tensor, Tuple[Tensor, Tensor]]: ... # type: ignore 80 | 81 | 82 | class GRU(RNNBase): 83 | def __init__(self, input_size: int, hidden_size: int, num_layers: int = ..., bias: bool = ..., 84 | batch_first: bool = ..., dropout: float = ..., bidirectional: bool = ..., 85 | nonlinearity: str = ...) -> None: ... 86 | 87 | def forward(self, input: Tensor, hx: Optional[Tensor] = ...) -> Tensor: ... # type: ignore 88 | 89 | def __call__(self, input: Tensor, hx: Optional[Tensor] = ...) -> Tensor: ... # type: ignore 90 | 91 | 92 | class RNNCellBase(Module): 93 | input_size: int = ... 94 | hidden_size: int = ... 95 | bias: bool = ... 96 | weight_ih: Parameter = ... 97 | weight_hh: Parameter = ... 98 | bias_ih: Parameter = ... 99 | bias_hh: Parameter = ... 100 | 101 | def __init__(self, input_size: int, hidden_size: int, bias: bool, num_chunks: int) -> None: ... 102 | 103 | def check_forward_input(self, input: Tensor) -> None: ... 104 | 105 | def check_forward_hidden(self, input: Tensor, hx: Tensor, hidden_label: str = ...) -> None: ... 106 | 107 | def reset_parameters(self) -> None: ... 108 | 109 | 110 | class RNNCell(RNNCellBase): 111 | nonlinearity: str = ... 112 | 113 | def __init__(self, input_size: int, hidden_size: int, bias: bool = ..., nonlinearity: str = ...) -> None: ... 114 | 115 | def forward(self, input: Tensor, hx: Optional[Tensor] = ...) -> Tensor: ... # type: ignore 116 | 117 | def __call__(self, input: Tensor, hx: Optional[Tensor] = ...) -> Tensor: ... # type: ignore 118 | 119 | 120 | class LSTMCell(RNNCellBase): 121 | def __init__(self, input_size: int, hidden_size: int, bias: bool = ...) -> None: ... 122 | 123 | def forward(self, input: Tensor, hx: Optional[Tuple[Tensor, Tensor]] = ...) -> Tuple[Tensor, Tensor]: ... # type: ignore 124 | 125 | def __call__(self, input: Tensor, hx: Optional[Tuple[Tensor, Tensor]] = ...) -> Tuple[Tensor, Tensor]: ... # type: ignore 126 | 127 | 128 | class GRUCell(RNNCellBase): 129 | def __init__(self, input_size: int, hidden_size: int, bias: bool = ...) -> None: ... 130 | 131 | def forward(self, input: Tensor, hx: Optional[Tensor] = ...) -> Tensor: ... # type: ignore 132 | 133 | def __call__(self, input: Tensor, hx: Optional[Tensor] = ...) -> Tensor: ... # type: ignore 134 | -------------------------------------------------------------------------------- /nn/modules/activation.pyi: -------------------------------------------------------------------------------- 1 | from ... import Tensor 2 | from .. import Parameter 3 | from .module import Module 4 | from typing import Any, Optional 5 | 6 | 7 | class Threshold(Module): 8 | threshold: float = ... 9 | value: float = ... 10 | inplace: bool = ... 11 | 12 | def __init__(self, threshold: float, value: float, inplace: bool = ...) -> None: ... 13 | 14 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 15 | 16 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 17 | 18 | 19 | class ReLU(Threshold): 20 | def __init__(self, inplace: bool = ...) -> None: ... 21 | 22 | 23 | class RReLU(Module): 24 | lower: float = ... 25 | upper: float = ... 26 | inplace: bool = ... 27 | 28 | def __init__(self, lower: float = ..., upper: float = ..., inplace: bool = ...) -> None: ... 29 | 30 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 31 | 32 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 33 | 34 | 35 | class Hardtanh(Module): 36 | min_val: float = ... 37 | max_val: float = ... 38 | inplace: bool = ... 39 | 40 | def __init__(self, min_val: float = ..., max_val: float = ..., inplace: bool = ...) -> None: ... 41 | 42 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 43 | 44 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 45 | 46 | 47 | class ReLU6(Hardtanh): 48 | def __init__(self, inplace: bool = ...) -> None: ... 49 | 50 | 51 | class Sigmoid(Module): 52 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 53 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 54 | 55 | 56 | class Tanh(Module): 57 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 58 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 59 | 60 | 61 | class ELU(Module): 62 | alpha: float = ... 63 | inplace: bool = ... 64 | 65 | def __init__(self, alpha: float = ..., inplace: bool = ...) -> None: ... 66 | 67 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 68 | 69 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 70 | 71 | 72 | class CELU(Module): 73 | alpha: float = ... 74 | inplace: bool = ... 75 | 76 | def __init__(self, alpha: float = ..., inplace: bool = ...) -> None: ... 77 | 78 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 79 | 80 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 81 | 82 | 83 | class SELU(Module): 84 | inplace: bool = ... 85 | 86 | def __init__(self, inplace: bool = ...) -> None: ... 87 | 88 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 89 | 90 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 91 | 92 | 93 | class GLU(Module): 94 | dim: int = ... 95 | 96 | def __init__(self, dim: int = ...) -> None: ... 97 | 98 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 99 | 100 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 101 | 102 | 103 | class GELU(Module): 104 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 105 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 106 | 107 | 108 | class Hardshrink(Module): 109 | lambd: float = ... 110 | 111 | def __init__(self, lambd: float = ...) -> None: ... 112 | 113 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 114 | 115 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 116 | 117 | 118 | class LeakyReLU(Module): 119 | negative_slope: float = ... 120 | inplace: bool = ... 121 | 122 | def __init__(self, negative_slope: float = ..., inplace: bool = ...) -> None: ... 123 | 124 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 125 | 126 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 127 | 128 | 129 | class LogSigmoid(Module): 130 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 131 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 132 | 133 | 134 | class Softplus(Module): 135 | beta: float = ... 136 | threshold: float = ... 137 | 138 | def __init__(self, beta: float = ..., threshold: float = ...) -> None: ... 139 | 140 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 141 | 142 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 143 | 144 | 145 | class Softshrink(Module): 146 | lambd: float = ... 147 | 148 | def __init__(self, lambd: float = ...) -> None: ... 149 | 150 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 151 | 152 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 153 | 154 | 155 | class PReLU(Module): 156 | num_parameters: int = ... 157 | weight: Parameter = ... 158 | 159 | def __init__(self, num_parameters: int = ..., init: float = ...) -> None: ... 160 | 161 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 162 | 163 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 164 | 165 | 166 | class Softsign(Module): 167 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 168 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 169 | 170 | 171 | class Tanhshrink(Module): 172 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 173 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 174 | 175 | 176 | class Softmin(Module): 177 | dim: int = ... 178 | 179 | def __init__(self, dim: Optional[int] = ...) -> None: ... 180 | 181 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 182 | 183 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 184 | 185 | 186 | class Softmax(Module): 187 | dim: int = ... 188 | 189 | def __init__(self, dim: Optional[int] = ...) -> None: ... 190 | 191 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 192 | 193 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 194 | 195 | 196 | class Softmax2d(Module): 197 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 198 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 199 | 200 | 201 | class LogSoftmax(Module): 202 | dim: int = ... 203 | 204 | def __init__(self, dim: Optional[int] = ...) -> None: ... 205 | 206 | def forward(self, input: Tensor) -> Tensor: ... # type: ignore 207 | 208 | def __call__(self, input: Tensor) -> Tensor: ... # type: ignore 209 | -------------------------------------------------------------------------------- /nn/modules/linear.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import torch 4 | from torch.nn.parameter import Parameter 5 | from .. import functional as F 6 | from .. import init 7 | from .module import Module 8 | 9 | 10 | class Identity(Module): 11 | r"""A placeholder identity operator that is argument-insensitive. 12 | 13 | Args: 14 | args: any argument (unused) 15 | kwargs: any keyword argument (unused) 16 | 17 | Examples:: 18 | 19 | >>> m = nn.Identity(54, unused_argument1=0.1, unused_argument2=False) 20 | >>> input = torch.randn(128, 20) 21 | >>> output = m(input) 22 | >>> print(output.size()) 23 | torch.Size([128, 20]) 24 | 25 | """ 26 | def __init__(self, *args, **kwargs): 27 | super(Identity, self).__init__() 28 | 29 | def forward(self, input): 30 | return input 31 | 32 | 33 | class Linear(Module): 34 | r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b` 35 | 36 | Args: 37 | in_features: size of each input sample 38 | out_features: size of each output sample 39 | bias: If set to ``False``, the layer will not learn an additive bias. 40 | Default: ``True`` 41 | 42 | Shape: 43 | - Input: :math:`(N, *, H_{in})` where :math:`*` means any number of 44 | additional dimensions and :math:`H_{in} = \text{in\_features}` 45 | - Output: :math:`(N, *, H_{out})` where all but the last dimension 46 | are the same shape as the input and :math:`H_{out} = \text{out\_features}`. 47 | 48 | Attributes: 49 | weight: the learnable weights of the module of shape 50 | :math:`(\text{out\_features}, \text{in\_features})`. The values are 51 | initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where 52 | :math:`k = \frac{1}{\text{in\_features}}` 53 | bias: the learnable bias of the module of shape :math:`(\text{out\_features})`. 54 | If :attr:`bias` is ``True``, the values are initialized from 55 | :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where 56 | :math:`k = \frac{1}{\text{in\_features}}` 57 | 58 | Examples:: 59 | 60 | >>> m = nn.Linear(20, 30) 61 | >>> input = torch.randn(128, 20) 62 | >>> output = m(input) 63 | >>> print(output.size()) 64 | torch.Size([128, 30]) 65 | """ 66 | __constants__ = ['bias', 'in_features', 'out_features'] 67 | 68 | def __init__(self, in_features, out_features, bias=True): 69 | super(Linear, self).__init__() 70 | self.in_features = in_features 71 | self.out_features = out_features 72 | self.weight = Parameter(torch.Tensor(out_features, in_features)) 73 | if bias: 74 | self.bias = Parameter(torch.Tensor(out_features)) 75 | else: 76 | self.register_parameter('bias', None) 77 | self.reset_parameters() 78 | 79 | def reset_parameters(self): 80 | init.kaiming_uniform_(self.weight, a=math.sqrt(5)) 81 | if self.bias is not None: 82 | fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) 83 | bound = 1 / math.sqrt(fan_in) 84 | init.uniform_(self.bias, -bound, bound) 85 | 86 | def forward(self, input): 87 | return F.linear(input, self.weight, self.bias) 88 | 89 | def extra_repr(self): 90 | return 'in_features={}, out_features={}, bias={}'.format( 91 | self.in_features, self.out_features, self.bias is not None 92 | ) 93 | 94 | 95 | class Bilinear(Module): 96 | r"""Applies a bilinear transformation to the incoming data: 97 | :math:`y = x_1 A x_2 + b` 98 | 99 | Args: 100 | in1_features: size of each first input sample 101 | in2_features: size of each second input sample 102 | out_features: size of each output sample 103 | bias: If set to False, the layer will not learn an additive bias. 104 | Default: ``True`` 105 | 106 | Shape: 107 | - Input1: :math:`(N, *, H_{in1})` where :math:`H_{in1}=\text{in1\_features}` and 108 | :math:`*` means any number of additional dimensions. All but the last dimension 109 | of the inputs should be the same. 110 | - Input2: :math:`(N, *, H_{in2})` where :math:`H_{in2}=\text{in2\_features}`. 111 | - Output: :math:`(N, *, H_{out})` where :math:`H_{out}=\text{out\_features}` 112 | and all but the last dimension are the same shape as the input. 113 | 114 | Attributes: 115 | weight: the learnable weights of the module of shape 116 | :math:`(\text{out\_features}, \text{in1\_features}, \text{in2\_features})`. 117 | The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where 118 | :math:`k = \frac{1}{\text{in1\_features}}` 119 | bias: the learnable bias of the module of shape :math:`(\text{out\_features})`. 120 | If :attr:`bias` is ``True``, the values are initialized from 121 | :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where 122 | :math:`k = \frac{1}{\text{in1\_features}}` 123 | 124 | Examples:: 125 | 126 | >>> m = nn.Bilinear(20, 30, 40) 127 | >>> input1 = torch.randn(128, 20) 128 | >>> input2 = torch.randn(128, 30) 129 | >>> output = m(input1, input2) 130 | >>> print(output.size()) 131 | torch.Size([128, 40]) 132 | """ 133 | __constants__ = ['in1_features', 'in2_features', 'out_features', 'bias'] 134 | 135 | def __init__(self, in1_features, in2_features, out_features, bias=True): 136 | super(Bilinear, self).__init__() 137 | self.in1_features = in1_features 138 | self.in2_features = in2_features 139 | self.out_features = out_features 140 | self.weight = Parameter(torch.Tensor(out_features, in1_features, in2_features)) 141 | 142 | if bias: 143 | self.bias = Parameter(torch.Tensor(out_features)) 144 | else: 145 | self.register_parameter('bias', None) 146 | self.reset_parameters() 147 | 148 | def reset_parameters(self): 149 | bound = 1 / math.sqrt(self.weight.size(1)) 150 | init.uniform_(self.weight, -bound, bound) 151 | if self.bias is not None: 152 | init.uniform_(self.bias, -bound, bound) 153 | 154 | def forward(self, input1, input2): 155 | return F.bilinear(input1, input2, self.weight, self.bias) 156 | 157 | def extra_repr(self): 158 | return 'in1_features={}, in2_features={}, out_features={}, bias={}'.format( 159 | self.in1_features, self.in2_features, self.out_features, self.bias is not None 160 | ) 161 | 162 | # TODO: PartialLinear - maybe in sparse? 163 | -------------------------------------------------------------------------------- /nn/quantized/modules/functional_modules.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch._ops import ops 3 | 4 | 5 | class FloatFunctional(torch.nn.Module): 6 | r"""State collector class for float operatitons. 7 | 8 | The instance of this class can be used instead of the ``torch.`` prefix for 9 | some operations. See example usage below. 10 | 11 | .. note:: 12 | 13 | This class does not provide a ``forward`` hook. Instead, you must use 14 | one of the underlying functions (e.g. ``add``). 15 | 16 | Examples:: 17 | 18 | >>> f_add = FloatFunctional() 19 | >>> a = torch.tensor(3.0) 20 | >>> b = torch.tensor(4.0) 21 | >>> f_add.add(a, b) # Equivalent to ``torch.add(3, 4) 22 | 23 | Valid operation names: 24 | - add 25 | - cat 26 | - mul 27 | - add_relu 28 | - add_scalar 29 | - mul_scalar 30 | """ 31 | def __init__(self): 32 | super(FloatFunctional, self).__init__() 33 | self.activation_post_process = torch.nn.Identity() 34 | 35 | def forward(self, x): 36 | raise RuntimeError("FloatFunctional is not intended to use the " + 37 | "'forward'. Please use the underlying operation") 38 | 39 | r"""Operation equivalent to ``torch.add(Tensor, Tensor)``""" 40 | def add(self, x, y): 41 | # type: (Tensor, Tensor) -> Tensor 42 | r = torch.add(x, y) 43 | r = self.activation_post_process(r) 44 | return r 45 | 46 | r"""Operation equivalent to ``torch.add(Tensor, float)``""" 47 | def add_scalar(self, x, y): 48 | # type: (Tensor, float) -> Tensor 49 | r = torch.add(x, y) 50 | # No observer needed for scalar add 51 | return r 52 | 53 | r"""Operation equivalent to ``torch.mul(Tensor, Tensor)``""" 54 | def mul(self, x, y): 55 | # type: (Tensor, Tensor) -> Tensor 56 | r = torch.mul(x, y) 57 | r = self.activation_post_process(r) 58 | return r 59 | 60 | r"""Operation equivalent to ``torch.mul(Tensor, float)``""" 61 | def mul_scalar(self, x, y): 62 | # type: (Tensor, float) -> Tensor 63 | r = torch.mul(x, y) 64 | # No observer needed for scalar multiply 65 | return r 66 | 67 | r"""Operation equivalent to ``torch.cat``""" 68 | def cat(self, x, dim=0): 69 | # type: (List[Tensor], int) -> Tensor 70 | r = torch.cat(x, dim=dim) 71 | r = self.activation_post_process(r) 72 | return r 73 | 74 | r"""Operation equivalent to ``relu(torch.add(x,y))``""" 75 | def add_relu(self, x, y): 76 | # type: (Tensor, Tensor) -> Tensor 77 | r = torch.add(x, y) 78 | r = torch.nn.functional.relu(r) 79 | r = self.activation_post_process(r) 80 | return r 81 | 82 | 83 | class QFunctional(torch.nn.Module): 84 | r"""Wrapper class for quantized operatitons. 85 | 86 | The instance of this class can be used instead of the 87 | ``torch.ops.quantized`` prefix. See example usage below. 88 | 89 | .. note:: 90 | 91 | This class does not provide a ``forward`` hook. Instead, you must use 92 | one of the underlying functions (e.g. ``add``). 93 | 94 | Examples:: 95 | 96 | >>> q_add = QFunctional('add') 97 | >>> a = torch.quantize_per_tensor(torch.tensor(3.0), 1.0, 0, torch.qint32) 98 | >>> b = torch.quantize_per_tensor(torch.tensor(4.0), 1.0, 0, torch.qint32) 99 | >>> q_add.add(a, b) # Equivalent to ``torch.ops.quantized.add(3, 4) 100 | 101 | Valid operation names: 102 | - add 103 | - cat 104 | - mul 105 | - add_relu 106 | - add_scalar 107 | - mul_scalar 108 | """ 109 | def __init__(self): 110 | super(QFunctional, self).__init__() 111 | self.scale = 1.0 112 | self.zero_point = 0 113 | 114 | def _save_to_state_dict(self, destination, prefix, keep_vars): 115 | super(QFunctional, self)._save_to_state_dict(destination, prefix, keep_vars) 116 | destination[prefix + 'scale'] = torch.tensor(self.scale) 117 | destination[prefix + 'zero_point'] = torch.tensor(self.zero_point) 118 | 119 | def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, 120 | missing_keys, unexpected_keys, error_msgs): 121 | 122 | self.scale = float(state_dict.pop(prefix + 'scale')) 123 | self.zero_point = int(state_dict.pop(prefix + 'zero_point')) 124 | super(QFunctional, self)._load_from_state_dict(state_dict, prefix, local_metadata, False, 125 | missing_keys, unexpected_keys, error_msgs) 126 | 127 | def forward(self, x): 128 | raise RuntimeError("Functional is not intended to use the " + 129 | "'forward'. Please use the underlying operation") 130 | 131 | r"""Operation equivalent to ``torch.ops.quantized.add``""" 132 | def add(self, x, y): 133 | # type: (Tensor, Tensor) -> Tensor 134 | return ops.quantized.add(x, y, scale=self.scale, 135 | zero_point=self.zero_point) 136 | 137 | r"""Operation equivalent to ``torch.ops.quantized.add(Tensor, float)``""" 138 | def add_scalar(self, x, y): 139 | # type: (Tensor, float) -> Tensor 140 | return ops.quantized.add_scalar(x, y) 141 | 142 | r"""Operation equivalent to ``torch.ops.quantized.mul(Tensor, Tensor)``""" 143 | def mul(self, x, y): 144 | # type: (Tensor, Tensor) -> Tensor 145 | return ops.quantized.mul(x, y, scale=self.scale, 146 | zero_point=self.zero_point) 147 | 148 | r"""Operation equivalent to ``torch.ops.quantized.mul(Tensor, float)``""" 149 | def mul_scalar(self, x, y): 150 | # type: (Tensor, float) -> Tensor 151 | return ops.quantized.mul_scalar(x, y) 152 | 153 | r"""Operation equivalent to ``torch.ops.quantized.cat``""" 154 | def cat(self, x, dim=0): 155 | # type: (List[Tensor], int) -> Tensor 156 | return ops.quantized.cat(x, scale=self.scale, 157 | zero_point=self.zero_point, dim=dim) 158 | 159 | r"""Operation equivalent to ``torch.ops.quantized.add_relu``""" 160 | def add_relu(self, x, y): 161 | # type: (Tensor, Tensor) -> Tensor 162 | return ops.quantized.add_relu(x, y, scale=self.scale, 163 | zero_point=self.zero_point) 164 | 165 | @classmethod 166 | def from_float(cls, mod): 167 | assert type(mod) == FloatFunctional,\ 168 | "QFunctional.from_float expects an instance of FloatFunctional" 169 | scale, zero_point = mod.activation_post_process.calculate_qparams() 170 | new_mod = QFunctional() 171 | new_mod.scale = float(scale) 172 | new_mod.zero_point = int(zero_point) 173 | return new_mod 174 | -------------------------------------------------------------------------------- /nn/parallel/replicate.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.cuda.comm as comm 3 | from torch.cuda._utils import _get_device_index 4 | from torch.nn import Parameter 5 | 6 | 7 | def _is_script_module(module): 8 | import torch.jit 9 | return isinstance(module, torch.jit.ScriptModule) 10 | 11 | 12 | def _is_script_method(module): 13 | import torch.jit 14 | return isinstance(module, torch._C.ScriptMethod) 15 | 16 | 17 | def _init_script_module(): 18 | import torch.jit 19 | return torch.jit.ScriptModule() 20 | 21 | 22 | def _is_jit_enabled(): 23 | import torch.jit 24 | return torch.jit._enabled 25 | 26 | 27 | # Check if we can safely replicate the module. 28 | # there are two types of module: 29 | # 1. python modules 30 | # 2. ScriptModule 31 | # 32 | # currently a module cannot be replicated properly if the descendants of 33 | # any ScriptModule contains python module (type 1 above) 34 | def _replicatable_module(module, memo=None): 35 | 36 | # module.modules() contains module itself as the first element 37 | def descendant_modules(module): 38 | gen = module.modules() 39 | next(gen) 40 | return gen 41 | 42 | if not _is_jit_enabled(): 43 | return True 44 | if memo is None: 45 | memo = set() 46 | 47 | # memorize visited modules 48 | memo.add(module) 49 | if _is_script_module(module): 50 | memo.update(descendant_modules(module)) 51 | return all(_is_script_module(descendant) for 52 | descendant in descendant_modules(module)) 53 | 54 | for child in module.children(): 55 | # since any unreplicatable module will cause the check to return 56 | # False early, visited modules here can be safely ignored. 57 | if child in memo: 58 | continue 59 | if not _replicatable_module(child, memo): 60 | return False 61 | 62 | return True 63 | 64 | def _broadcast_coalesced_reshape(tensors, devices, detach=False): 65 | from ._functions import Broadcast 66 | if detach: 67 | return comm.broadcast_coalesced(tensors, devices) 68 | else: 69 | # Use the autograd function to broadcast if not detach 70 | if len(tensors) > 0: 71 | tensor_copies = Broadcast.apply(devices, *tensors) 72 | return [tensor_copies[i:i + len(tensors)] 73 | for i in range(0, len(tensor_copies), len(tensors))] 74 | else: 75 | return [] 76 | 77 | 78 | def replicate(network, devices, detach=False): 79 | if not _replicatable_module(network): 80 | raise RuntimeError("Cannot replicate network where python modules are " 81 | "childrens of ScriptModule") 82 | 83 | devices = list(map(lambda x: _get_device_index(x, True), devices)) 84 | num_replicas = len(devices) 85 | 86 | params = list(network.parameters()) 87 | param_indices = {param: idx for idx, param in enumerate(params)} 88 | param_copies = _broadcast_coalesced_reshape(params, devices, detach) 89 | 90 | buffers = list(network.buffers()) 91 | buffers_rg = [] 92 | buffers_not_rg = [] 93 | for buf in buffers: 94 | if buf.requires_grad and not detach: 95 | buffers_rg.append(buf) 96 | else: 97 | buffers_not_rg.append(buf) 98 | 99 | buffer_indices_rg = {buf: idx for idx, buf in enumerate(buffers_rg)} 100 | buffer_indices_not_rg = {buf: idx for idx, buf in enumerate(buffers_not_rg)} 101 | 102 | buffer_copies_rg = _broadcast_coalesced_reshape(buffers_rg, devices, detach=detach) 103 | buffer_copies_not_rg = _broadcast_coalesced_reshape(buffers_not_rg, devices, detach=True) 104 | 105 | modules = list(network.modules()) 106 | module_copies = [[] for device in devices] 107 | module_indices = {} 108 | scriptmodule_skip_attr = {"_parameters", "_buffers", "_modules", "forward", "_c"} 109 | 110 | for i, module in enumerate(modules): 111 | module_indices[module] = i 112 | for j in range(num_replicas): 113 | if _is_script_module(module): 114 | # we have to initialize ScriptModule properly so that 115 | # it works with pybind11 116 | def init_fn(script_module): 117 | # Don't do anything here, we'll initialize the ScriptModule below 118 | return 119 | replica = torch.jit.RecursiveScriptModule._construct(module._c._replicate_for_data_parallel(), init_fn) 120 | else: 121 | replica = module._replicate_for_data_parallel() 122 | 123 | module_copies[j].append(replica) 124 | 125 | for i, module in enumerate(modules): 126 | for key, child in module._modules.items(): 127 | if child is None: 128 | for j in range(num_replicas): 129 | replica = module_copies[j][i] 130 | replica._modules[key] = None 131 | else: 132 | module_idx = module_indices[child] 133 | for j in range(num_replicas): 134 | replica = module_copies[j][i] 135 | setattr(replica, key, module_copies[j][module_idx]) 136 | for key, param in module._parameters.items(): 137 | if param is None: 138 | for j in range(num_replicas): 139 | replica = module_copies[j][i] 140 | replica._parameters[key] = None 141 | else: 142 | param_idx = param_indices[param] 143 | for j in range(num_replicas): 144 | replica = module_copies[j][i] 145 | param = param_copies[j][param_idx] 146 | setattr(replica, key, Parameter(param)) 147 | # TODO: We need to manually set _parameters with a bare 148 | # non-parameter Tensor, otherwise gradients don't 149 | # accumulate in the original parameters when you call 150 | # backwards() on the DataParallel module. 151 | replica._parameters[key] = param 152 | for key, buf in module._buffers.items(): 153 | if buf is None: 154 | for j in range(num_replicas): 155 | replica = module_copies[j][i] 156 | replica._buffers[key] = None 157 | else: 158 | if buf.requires_grad and not detach: 159 | buffer_copies = buffer_copies_rg 160 | buffer_idx = buffer_indices_rg[buf] 161 | else: 162 | buffer_copies = buffer_copies_not_rg 163 | buffer_idx = buffer_indices_not_rg[buf] 164 | for j in range(num_replicas): 165 | replica = module_copies[j][i] 166 | setattr(replica, key, buffer_copies[j][buffer_idx]) 167 | 168 | return [module_copies[j][0] for j in range(num_replicas)] 169 | -------------------------------------------------------------------------------- /quantization/fuse_modules.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | 3 | import torch 4 | import copy 5 | 6 | import torch.nn.intrinsic.modules.fused as torch_fused 7 | 8 | def fuse_conv_bn(conv, bn): 9 | r"""Given the conv and bn modules, fuses them and returns the fused module 10 | 11 | Args: 12 | conv: Module instance of type conv2d 13 | bn: Spatial BN instance that needs to be fused with the conv 14 | 15 | Examples:: 16 | 17 | >>> m1 = nn.Conv2d(10, 20, 3) 18 | >>> b1 = nn.BatchNorm2d(20) 19 | >>> m2 = fuse_conv_bn(m1, b1) 20 | """ 21 | assert(conv.training == bn.training),\ 22 | "Conv and BN both must be in the same mode (train or eval)." 23 | 24 | if conv.training: 25 | assert conv.bias is None, 'Only support fusing Conv2d that does not have bias' 26 | assert bn.num_features == conv.out_channels, 'Output channel of Conv2d must match num_features of BatchNorm2d' 27 | assert bn.affine, 'Only support fusing BatchNorm2d with affine set to True' 28 | assert bn.track_running_stats, 'Only support fusing BatchNorm2d with tracking_running_stats set to True' 29 | return torch.nn.intrinsic.ConvBn2d(conv, bn) 30 | else: 31 | return torch.nn.utils.fuse_conv_bn_eval(conv, bn) 32 | 33 | def fuse_conv_bn_relu(conv, bn, relu): 34 | r"""Given the conv and bn modules, fuses them and returns the fused module 35 | 36 | Args: 37 | conv: Module instance of type conv2d 38 | bn: Spatial BN instance that needs to be fused with the conv 39 | 40 | Examples:: 41 | 42 | >>> m1 = nn.Conv2d(10, 20, 3) 43 | >>> b1 = nn.BatchNorm2d(20) 44 | >>> m2 = fuse_conv_bn(m1, b1) 45 | """ 46 | assert(conv.training == bn.training == relu.training),\ 47 | "Conv and BN both must be in the same mode (train or eval)." 48 | 49 | if conv.training: 50 | assert not relu.inplace, 'We only support fusion of non-inplace ReLU.' 51 | return torch_fused.ConvBnReLU2d(conv, bn, relu) 52 | else: 53 | return torch_fused.ConvReLU2d( 54 | torch.nn.utils.fusion.fuse_conv_bn_eval(conv, bn), relu) 55 | 56 | # Generalization of getattr 57 | def _get_module(model, submodule_key): 58 | tokens = submodule_key.split('.') 59 | cur_mod = model 60 | for s in tokens: 61 | cur_mod = getattr(cur_mod, s) 62 | return cur_mod 63 | 64 | # Generalization of setattr 65 | def _set_module(model, submodule_key, module): 66 | tokens = submodule_key.split('.') 67 | sub_tokens = tokens[:-1] 68 | cur_mod = model 69 | for s in sub_tokens: 70 | cur_mod = getattr(cur_mod, s) 71 | 72 | setattr(cur_mod, tokens[-1], module) 73 | 74 | def fuse_known_modules(mod_list): 75 | r"""Returns a list of modules that fuses the operations specified 76 | in the input module list. 77 | 78 | Fuses only the following sequence of modules: 79 | conv, bn 80 | conv, bn, relu 81 | conv, relu 82 | linear, relu 83 | For these sequences, the first element in the output module list performs 84 | the fused operation. The rest of the elements are set to nn.Identity() 85 | """ 86 | 87 | OP_LIST_TO_FUSER_METHOD = { 88 | (torch.nn.Conv2d, torch.nn.BatchNorm2d): fuse_conv_bn, 89 | (torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.ReLU): fuse_conv_bn_relu, 90 | (torch.nn.Conv2d, torch.nn.ReLU): torch.nn.intrinsic.ConvReLU2d, 91 | (torch.nn.Linear, torch.nn.ReLU): torch.nn.intrinsic.LinearReLU 92 | } 93 | 94 | types = tuple(type(m) for m in mod_list) 95 | fuser_method = OP_LIST_TO_FUSER_METHOD.get(types, None) 96 | if fuser_method is None: 97 | raise NotImplementedError("Cannot fuse modules: {}".format(types)) 98 | new_mod = [None] * len(mod_list) 99 | new_mod[0] = fuser_method(*mod_list) 100 | 101 | for i in range(1, len(mod_list)): 102 | new_mod[i] = torch.nn.Identity() 103 | new_mod[i].training = mod_list[0].training 104 | 105 | return new_mod 106 | 107 | def _fuse_modules(model, modules_to_fuse, fuser_func=fuse_known_modules): 108 | 109 | mod_list = [] 110 | for item in modules_to_fuse: 111 | mod_list.append(_get_module(model, item)) 112 | 113 | # Fuse list of modules 114 | new_mod_list = fuser_func(mod_list) 115 | 116 | # Replace original module list with fused module list 117 | for i, item in enumerate(modules_to_fuse): 118 | _set_module(model, item, new_mod_list[i]) 119 | 120 | def fuse_modules(model, modules_to_fuse, inplace=False, fuser_func=fuse_known_modules): 121 | r"""Fuses a list of modules into a single module 122 | 123 | Fuses only the following sequence of modules: 124 | 125 | * conv, bn 126 | 127 | * conv, bn, relu 128 | 129 | * conv, relu 130 | 131 | * linear, relu 132 | 133 | All other sequences are left unchanged. 134 | For these sequences, replaces the first item in the list 135 | with the fused module, replacing the rest of the modules 136 | with identity. 137 | 138 | Arguments: 139 | model: Model containing the modules to be fused 140 | modules_to_fuse: list of list of module names to fuse. Can also be a list 141 | of strings if there is only a single list of modules to fuse. 142 | inplace: bool specifying if fusion happens in place on the model, by default 143 | a new model is returned 144 | fuser_func: Function that takes in a list of modules and outputs a list of fused modules 145 | of the same length. For example, 146 | fuser_func([convModule, BNModule]) returns the list [ConvBNModule, nn.Identity()] 147 | Defaults to torch.quantization.fuse_known_modules 148 | Returns: 149 | model with fused modules. A new copy is created if inplace=True. 150 | 151 | Examples:: 152 | 153 | >>> m = myModel() 154 | >>> # m is a module containing the sub-modules below 155 | >>> modules_to_fuse = [ ['conv1', 'bn1', 'relu1'], ['submodule.conv', 'submodule.relu']] 156 | >>> fused_m = torch.quantization.fuse_modules(m, modules_to_fuse) 157 | >>> output = fused_m(input) 158 | 159 | >>> m = myModel() 160 | >>> # Alternately provide a single list of modules to fuse 161 | >>> modules_to_fuse = ['conv1', 'bn1', 'relu1'] 162 | >>> fused_m = torch.quantization.fuse_modules(m, modules_to_fuse) 163 | >>> output = fused_m(input) 164 | 165 | """ 166 | if not inplace: 167 | model = copy.deepcopy(model) 168 | 169 | if all(isinstance(module_element, str) for module_element in modules_to_fuse): 170 | # Handle case of modules_to_fuse being a list 171 | _fuse_modules(model, modules_to_fuse, fuser_func) 172 | else: 173 | # Handle case of modules_to_fuse being a list of lists 174 | for module_list in modules_to_fuse: 175 | _fuse_modules(model, module_list, fuser_func) 176 | return model 177 | -------------------------------------------------------------------------------- /nn/modules/loss.pyi: -------------------------------------------------------------------------------- 1 | from typing import Any, Optional 2 | from .module import Module 3 | from ... import Tensor 4 | 5 | 6 | # The deprecated `size_average` and `reduce` arguments are not included in the stubs 7 | class _Loss(Module): 8 | reduction: str = ... 9 | 10 | def __init__(self, reduction: str = ...) -> None: ... 11 | 12 | 13 | class _WeightedLoss(_Loss): 14 | def __init__(self, weight: Optional[Any] = ..., reduction: str = ...) -> None: ... 15 | 16 | 17 | class L1Loss(_Loss): 18 | def __init__(self, reduction: str = ...) -> None: ... 19 | 20 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 21 | 22 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 23 | 24 | 25 | class NLLLoss(_WeightedLoss): 26 | ignore_index: int = ... 27 | 28 | def __init__(self, weight: Optional[Any] = ..., ignore_index: int = ..., reduction: str = ...) -> None: ... 29 | 30 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 31 | 32 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 33 | 34 | 35 | class NLLLoss2d(NLLLoss): 36 | def __init__(self, weight: Optional[Any] = ..., ignore_index: int = ..., reduction: str = ...) -> None: ... 37 | 38 | 39 | class PoissonNLLLoss(_Loss): 40 | log_input: bool = ... 41 | full: bool = ... 42 | eps: float = ... 43 | 44 | def __init__(self, log_input: bool = ..., full: bool = ..., eps: float = ..., reduction: str = ...) -> None: ... 45 | 46 | def forward(self, log_input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 47 | 48 | def __call__(self, log_input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 49 | 50 | 51 | class KLDivLoss(_Loss): 52 | def __init__(self, reduction: str = ...) -> None: ... 53 | 54 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 55 | 56 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 57 | 58 | 59 | class MSELoss(_Loss): 60 | def __init__(self, reduction: str = ...) -> None: ... 61 | 62 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 63 | 64 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 65 | 66 | 67 | class BCELoss(_WeightedLoss): 68 | def __init__(self, weight: Optional[Any] = ..., reduction: str = ...) -> None: ... 69 | 70 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 71 | 72 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 73 | 74 | 75 | class BCEWithLogitsLoss(_Loss): 76 | def __init__(self, weight: Optional[Any] = ..., reduction: str = ..., pos_weight: Optional[Any] = ...) -> None: ... 77 | 78 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 79 | 80 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 81 | 82 | 83 | class HingeEmbeddingLoss(_Loss): 84 | margin: Any = ... 85 | 86 | def __init__(self, margin: float = ..., reduction: str = ...) -> None: ... 87 | 88 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 89 | 90 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 91 | 92 | 93 | class MultiLabelMarginLoss(_Loss): 94 | def __init__(self, reduction: str = ...) -> None: ... 95 | 96 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 97 | 98 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 99 | 100 | 101 | class SmoothL1Loss(_Loss): 102 | def __init__(self, reduction: str = ...) -> None: ... 103 | 104 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 105 | 106 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 107 | 108 | 109 | class SoftMarginLoss(_Loss): 110 | def __init__(self, reduction: str = ...) -> None: ... 111 | 112 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 113 | 114 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 115 | 116 | 117 | class CrossEntropyLoss(_WeightedLoss): 118 | ignore_index: int = ... 119 | 120 | def __init__(self, weight: Optional[Any] = ..., ignore_index: int = ..., reduction: str = ...) -> None: ... 121 | 122 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 123 | 124 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 125 | 126 | 127 | class MultiLabelSoftMarginLoss(_WeightedLoss): 128 | def __init__(self, weight: Optional[Any] = ..., reduction: str = ...) -> None: ... 129 | 130 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 131 | 132 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 133 | 134 | 135 | class CosineEmbeddingLoss(_Loss): 136 | margin: float = ... 137 | 138 | def __init__(self, margin: float = ..., reduction: str = ...) -> None: ... 139 | 140 | def forward(self, input1: Tensor, input2: Tensor, target: Tensor) -> Tensor: ... # type: ignore 141 | 142 | def __call__(self, input1: Tensor, input2: Tensor, target: Tensor) -> Tensor: ... # type: ignore 143 | 144 | 145 | class MarginRankingLoss(_Loss): 146 | margin: float = ... 147 | 148 | def __init__(self, margin: float = ..., reduction: str = ...) -> None: ... 149 | 150 | def forward(self, input1: Tensor, input2: Tensor, target: Tensor) -> Tensor: ... # type: ignore 151 | 152 | def __call__(self, input1: Tensor, input2: Tensor, target: Tensor) -> Tensor: ... # type: ignore 153 | 154 | 155 | class MultiMarginLoss(_WeightedLoss): 156 | p: int = ... 157 | margin: float = ... 158 | 159 | def __init__(self, p: int = ..., margin: float = ..., weight: Optional[Any] = ..., 160 | reduction: str = ...) -> None: ... 161 | 162 | def forward(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 163 | 164 | def __call__(self, input: Tensor, target: Tensor) -> Tensor: ... # type: ignore 165 | 166 | 167 | class TripletMarginLoss(_Loss): 168 | margin: float = ... 169 | p: int = ... 170 | eps: float = ... 171 | swap: bool = ... 172 | 173 | def __init__(self, margin: float = ..., p: int = ..., eps: float = ..., swap: bool = ..., 174 | reduction: str = ...) -> None: ... 175 | 176 | def forward(self, anchor: Tensor, positive: Tensor, negative: Tensor) -> Tensor: ... # type: ignore 177 | 178 | def __call__(self, anchor: Tensor, positive: Tensor, negative: Tensor) -> Tensor: ... # type: ignore 179 | 180 | 181 | class CTCLoss(_Loss): 182 | blank: int = ... 183 | zero_infinity: bool = ... 184 | 185 | def __init__(self, blank: int = ..., reduction: str = ..., zero_infinity: bool = ...) -> None: ... 186 | 187 | def forward(self, log_probs: Tensor, targets: Tensor, input_lengths: Tensor, target_lengths: Tensor) -> Tensor: ... # type: ignore 188 | 189 | def __call__(self, log_probs: Tensor, targets: Tensor, input_lengths: Tensor, target_lengths: Tensor) -> Tensor: ... # type: ignore 190 | -------------------------------------------------------------------------------- /nn/modules/dropout.py: -------------------------------------------------------------------------------- 1 | from .module import Module 2 | from .. import functional as F 3 | 4 | 5 | class _DropoutNd(Module): 6 | __constants__ = ['p', 'inplace'] 7 | 8 | def __init__(self, p=0.5, inplace=False): 9 | super(_DropoutNd, self).__init__() 10 | if p < 0 or p > 1: 11 | raise ValueError("dropout probability has to be between 0 and 1, " 12 | "but got {}".format(p)) 13 | self.p = p 14 | self.inplace = inplace 15 | 16 | def extra_repr(self): 17 | return 'p={}, inplace={}'.format(self.p, self.inplace) 18 | 19 | 20 | class Dropout(_DropoutNd): 21 | r"""During training, randomly zeroes some of the elements of the input 22 | tensor with probability :attr:`p` using samples from a Bernoulli 23 | distribution. Each channel will be zeroed out independently on every forward 24 | call. 25 | 26 | This has proven to be an effective technique for regularization and 27 | preventing the co-adaptation of neurons as described in the paper 28 | `Improving neural networks by preventing co-adaptation of feature 29 | detectors`_ . 30 | 31 | Furthermore, the outputs are scaled by a factor of :math:`\frac{1}{1-p}` during 32 | training. This means that during evaluation the module simply computes an 33 | identity function. 34 | 35 | Args: 36 | p: probability of an element to be zeroed. Default: 0.5 37 | inplace: If set to ``True``, will do this operation in-place. Default: ``False`` 38 | 39 | Shape: 40 | - Input: :math:`(*)`. Input can be of any shape 41 | - Output: :math:`(*)`. Output is of the same shape as input 42 | 43 | Examples:: 44 | 45 | >>> m = nn.Dropout(p=0.2) 46 | >>> input = torch.randn(20, 16) 47 | >>> output = m(input) 48 | 49 | .. _Improving neural networks by preventing co-adaptation of feature 50 | detectors: https://arxiv.org/abs/1207.0580 51 | """ 52 | 53 | def forward(self, input): 54 | return F.dropout(input, self.p, self.training, self.inplace) 55 | 56 | 57 | class Dropout2d(_DropoutNd): 58 | r"""Randomly zero out entire channels (a channel is a 2D feature map, 59 | e.g., the :math:`j`-th channel of the :math:`i`-th sample in the 60 | batched input is a 2D tensor :math:`\text{input}[i, j]`). 61 | Each channel will be zeroed out independently on every forward call with 62 | probability :attr:`p` using samples from a Bernoulli distribution. 63 | 64 | Usually the input comes from :class:`nn.Conv2d` modules. 65 | 66 | As described in the paper 67 | `Efficient Object Localization Using Convolutional Networks`_ , 68 | if adjacent pixels within feature maps are strongly correlated 69 | (as is normally the case in early convolution layers) then i.i.d. dropout 70 | will not regularize the activations and will otherwise just result 71 | in an effective learning rate decrease. 72 | 73 | In this case, :func:`nn.Dropout2d` will help promote independence between 74 | feature maps and should be used instead. 75 | 76 | Args: 77 | p (float, optional): probability of an element to be zero-ed. 78 | inplace (bool, optional): If set to ``True``, will do this operation 79 | in-place 80 | 81 | Shape: 82 | - Input: :math:`(N, C, H, W)` 83 | - Output: :math:`(N, C, H, W)` (same shape as input) 84 | 85 | Examples:: 86 | 87 | >>> m = nn.Dropout2d(p=0.2) 88 | >>> input = torch.randn(20, 16, 32, 32) 89 | >>> output = m(input) 90 | 91 | .. _Efficient Object Localization Using Convolutional Networks: 92 | http://arxiv.org/abs/1411.4280 93 | """ 94 | 95 | def forward(self, input): 96 | return F.dropout2d(input, self.p, self.training, self.inplace) 97 | 98 | 99 | class Dropout3d(_DropoutNd): 100 | r"""Randomly zero out entire channels (a channel is a 3D feature map, 101 | e.g., the :math:`j`-th channel of the :math:`i`-th sample in the 102 | batched input is a 3D tensor :math:`\text{input}[i, j]`). 103 | Each channel will be zeroed out independently on every forward call with 104 | probability :attr:`p` using samples from a Bernoulli distribution. 105 | 106 | Usually the input comes from :class:`nn.Conv3d` modules. 107 | 108 | As described in the paper 109 | `Efficient Object Localization Using Convolutional Networks`_ , 110 | if adjacent pixels within feature maps are strongly correlated 111 | (as is normally the case in early convolution layers) then i.i.d. dropout 112 | will not regularize the activations and will otherwise just result 113 | in an effective learning rate decrease. 114 | 115 | In this case, :func:`nn.Dropout3d` will help promote independence between 116 | feature maps and should be used instead. 117 | 118 | Args: 119 | p (float, optional): probability of an element to be zeroed. 120 | inplace (bool, optional): If set to ``True``, will do this operation 121 | in-place 122 | 123 | Shape: 124 | - Input: :math:`(N, C, D, H, W)` 125 | - Output: :math:`(N, C, D, H, W)` (same shape as input) 126 | 127 | Examples:: 128 | 129 | >>> m = nn.Dropout3d(p=0.2) 130 | >>> input = torch.randn(20, 16, 4, 32, 32) 131 | >>> output = m(input) 132 | 133 | .. _Efficient Object Localization Using Convolutional Networks: 134 | http://arxiv.org/abs/1411.4280 135 | """ 136 | 137 | def forward(self, input): 138 | return F.dropout3d(input, self.p, self.training, self.inplace) 139 | 140 | 141 | class AlphaDropout(_DropoutNd): 142 | r"""Applies Alpha Dropout over the input. 143 | 144 | Alpha Dropout is a type of Dropout that maintains the self-normalizing 145 | property. 146 | For an input with zero mean and unit standard deviation, the output of 147 | Alpha Dropout maintains the original mean and standard deviation of the 148 | input. 149 | Alpha Dropout goes hand-in-hand with SELU activation function, which ensures 150 | that the outputs have zero mean and unit standard deviation. 151 | 152 | During training, it randomly masks some of the elements of the input 153 | tensor with probability *p* using samples from a bernoulli distribution. 154 | The elements to masked are randomized on every forward call, and scaled 155 | and shifted to maintain zero mean and unit standard deviation. 156 | 157 | During evaluation the module simply computes an identity function. 158 | 159 | More details can be found in the paper `Self-Normalizing Neural Networks`_ . 160 | 161 | Args: 162 | p (float): probability of an element to be dropped. Default: 0.5 163 | inplace (bool, optional): If set to ``True``, will do this operation 164 | in-place 165 | 166 | Shape: 167 | - Input: :math:`(*)`. Input can be of any shape 168 | - Output: :math:`(*)`. Output is of the same shape as input 169 | 170 | Examples:: 171 | 172 | >>> m = nn.AlphaDropout(p=0.2) 173 | >>> input = torch.randn(20, 16) 174 | >>> output = m(input) 175 | 176 | .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515 177 | """ 178 | 179 | def forward(self, input): 180 | return F.alpha_dropout(input, self.p, self.training) 181 | 182 | 183 | class FeatureAlphaDropout(_DropoutNd): 184 | 185 | def forward(self, input): 186 | return F.feature_alpha_dropout(input, self.p, self.training) 187 | -------------------------------------------------------------------------------- /quantization/fake_quantize.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | import torch 3 | from torch.nn import Module 4 | from .observer import MovingAverageMinMaxObserver, HistogramObserver, MovingAveragePerChannelMinMaxObserver, _with_args 5 | 6 | class FakeQuantize(Module): 7 | r""" Simulate the quantize and dequantize operations in training time. 8 | The output of this module is given by 9 | 10 | x_out = (clamp(round(x/scale + zero_point), quant_min, quant_max)-zero_point)*scale 11 | 12 | 13 | 14 | * :attr:`scale` defines the scale factor used for quantization. 15 | 16 | * :attr:`zero_point` specifies the quantized value to which 0 in floating point maps to 17 | 18 | * :attr:`quant_min` specifies the minimum allowable quantized value. 19 | 20 | * :attr:`quant_max` specifies the maximum allowable quantized value. 21 | 22 | * :attr:`fake_quant_enable` controls the application of fake quantization on tensors, note that 23 | statistics can still be updated. 24 | 25 | * :attr:`observer_enable` controls statistics collection on tensors 26 | 27 | * :attr:`dtype` specifies the quantized dtype that is being emulated with fake-quantization, 28 | allowable values are torch.qint8 and torch.quint8. The values of quant_min and 29 | quant_max should be chosen to be consistent with the dtype 30 | 31 | 32 | Args: 33 | observer (module): Module for observing statistics on input tensors and calculating scale 34 | and zero-point. 35 | quant_min (int): The minimum allowable quantized value. 36 | quant_max (int): The maximum allowable quantized value. 37 | observer_kwargs (optional): Arguments for the observer module 38 | 39 | Attributes: 40 | observer (Module): User provided module that collects statistics on the input tensor and 41 | provides a method to calculate scale and zero-point. 42 | 43 | """ 44 | def __init__(self, observer=MovingAverageMinMaxObserver, quant_min=0, quant_max=255, **observer_kwargs): 45 | super(FakeQuantize, self).__init__() 46 | assert quant_min <= quant_max, \ 47 | 'quant_min must be less than or equal to quant_max' 48 | self.quant_min = quant_min 49 | self.quant_max = quant_max 50 | self.fake_quant_enabled = True 51 | self.observer_enabled = True 52 | self.activation_post_process = observer(**observer_kwargs) 53 | assert torch.iinfo(self.activation_post_process.dtype).min <= quant_min, 'quant_min out of bound' 54 | assert quant_max <= torch.iinfo(self.activation_post_process.dtype).max, 'quant_max out of bound' 55 | self.scale = None 56 | self.zero_point = None 57 | self.dtype = self.activation_post_process.dtype 58 | self.qscheme = self.activation_post_process.qscheme 59 | self.ch_axis = self.activation_post_process.ch_axis if hasattr(self.activation_post_process, 'ch_axis') else 0 60 | 61 | def enable_fake_quant(self, enabled=True): 62 | self.fake_quant_enabled = enabled 63 | return self 64 | 65 | def disable_fake_quant(self): 66 | return self.enable_fake_quant(False) 67 | 68 | def enable_observer(self, enabled=True): 69 | self.observer_enabled = enabled 70 | return self 71 | 72 | def disable_observer(self): 73 | return self.enable_observer(False) 74 | 75 | def calculate_qparams(self): 76 | return self.activation_post_process.calculate_qparams() 77 | 78 | def forward(self, X): 79 | if self.observer_enabled: 80 | self.activation_post_process(X.detach()) 81 | qparams = self.calculate_qparams() 82 | self.scale, self.zero_point = qparams[0], qparams[1] 83 | if self.fake_quant_enabled: 84 | if self.qscheme == torch.per_channel_symmetric or self.qscheme == torch.per_channel_affine: 85 | X = torch.fake_quantize_per_channel_affine(X, self.scale, self.zero_point, 86 | self.ch_axis, self.quant_min, self.quant_max) 87 | else: 88 | X = torch.fake_quantize_per_tensor_affine(X, float(self.scale), 89 | int(self.zero_point), self.quant_min, 90 | self.quant_max) 91 | return X 92 | 93 | with_args = classmethod(_with_args) 94 | 95 | def extra_repr(self): 96 | return 'fake_quant_enabled={}, observer_enabled={},\ 97 | scale={}, zero_point={}'.format( 98 | self.fake_quant_enabled, self.observer_enabled, 99 | self.scale, self.zero_point) 100 | 101 | def _save_to_state_dict(self, destination, prefix, keep_vars): 102 | # We cannot currently register scalar values as buffers, so need to manually 103 | # specify serialization here. 104 | super(FakeQuantize, self)._save_to_state_dict(destination, prefix, keep_vars) 105 | destination[prefix + 'scale'] = self.scale 106 | destination[prefix + 'zero_point'] = self.zero_point 107 | 108 | def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, 109 | missing_keys, unexpected_keys, error_msgs): 110 | 111 | self.scale = state_dict.pop(prefix + 'scale') 112 | self.zero_point = state_dict.pop(prefix + 'zero_point') 113 | super(FakeQuantize, self)._load_from_state_dict(state_dict, prefix, local_metadata, False, 114 | missing_keys, unexpected_keys, error_msgs) 115 | 116 | default_fake_quant = FakeQuantize.with_args(observer=MovingAverageMinMaxObserver, quant_min=0, quant_max=255, 117 | dtype=torch.quint8, qscheme=torch.per_tensor_affine, reduce_range=True) 118 | default_weight_fake_quant = FakeQuantize.with_args(observer=MovingAverageMinMaxObserver, quant_min=-128, quant_max=127, 119 | dtype=torch.qint8, qscheme=torch.per_tensor_symmetric, reduce_range=False) 120 | 121 | default_per_channel_weight_fake_quant = FakeQuantize.with_args(observer=MovingAveragePerChannelMinMaxObserver, 122 | quant_min=-128, 123 | quant_max=127, 124 | dtype=torch.qint8, 125 | qscheme=torch.per_channel_symmetric, 126 | reduce_range=False, 127 | ch_axis=0) 128 | default_histogram_fake_quant = FakeQuantize.with_args(observer=HistogramObserver, 129 | quant_min=0, 130 | quant_max=255, 131 | dtype=torch.quint8, 132 | qscheme=torch.per_tensor_affine, 133 | reduce_range=True) 134 | def disable_fake_quant(mod): 135 | if type(mod) == FakeQuantize: 136 | mod.disable_fake_quant() 137 | 138 | def enable_fake_quant(mod): 139 | if type(mod) == FakeQuantize: 140 | mod.enable_fake_quant() 141 | 142 | def disable_observer(mod): 143 | if type(mod) == FakeQuantize: 144 | mod.disable_observer() 145 | 146 | def enable_observer(mod): 147 | if type(mod) == FakeQuantize: 148 | mod.enable_observer() 149 | --------------------------------------------------------------------------------