├── .gitattributes ├── neuralnet_pytorch ├── zoo │ ├── __init__.py │ ├── vgg.py │ └── resnet.py ├── optim │ ├── lr_scheduler │ │ ├── __init__.py │ │ ├── inverse_lr.py │ │ └── warm_restart.py │ ├── __init__.py │ ├── nadam.py │ ├── lookahead.py │ └── adabound.py ├── extensions │ ├── pc2vox.py │ ├── bpd.py │ ├── include │ │ ├── bpd.h │ │ ├── pc2vox.h │ │ ├── chamfer_cuda.h │ │ ├── emd_cuda.h │ │ └── utils.h │ ├── __init__.py │ ├── csrc │ │ ├── bpd.cpp │ │ ├── utils.cpp │ │ ├── bindings.cpp │ │ ├── chamfer_cuda.cpp │ │ ├── emd_cuda.cpp │ │ ├── pc2vox.cpp │ │ ├── chamfer_kernel.cu │ │ └── emd_kernel.cu │ ├── dist_emd.py │ └── dist_chamfer.py ├── layers │ ├── __init__.py │ ├── adain.py │ ├── aggregation.py │ └── points.py ├── gin_nnt │ ├── __init__.py │ └── external_configurables.py ├── utils │ ├── __init__.py │ ├── misc_utils.py │ ├── activation_utils.py │ ├── numpy_utils.py │ ├── layer_utils.py │ └── cv_utils.py ├── version.py ├── __init__.py └── metrics.py ├── docs ├── source │ ├── _static │ │ └── nnt-logo.png │ ├── manual │ │ ├── metrics.rst │ │ ├── monitor.rst │ │ ├── track.rst │ │ ├── optimization.rst │ │ ├── zoo.rst │ │ ├── resizing.rst │ │ ├── normalization.rst │ │ ├── index.rst │ │ ├── utilities.rst │ │ ├── layers.rst │ │ └── gin.rst │ ├── help.rst │ ├── license.rst │ ├── index.rst │ ├── installation.rst │ └── conf.py ├── Makefile └── make.bat ├── .gitignore ├── tests ├── test_files │ └── test.gin ├── misc_test.py └── function_test.py ├── MANIFEST.in ├── requirements.txt ├── setup.cfg ├── .travis ├── travis_install.sh └── travis_before_install.sh ├── examples ├── configs │ └── cifar.gin └── cifar.py ├── LICENSE.txt ├── .travis.yml ├── .clang-format ├── setup.py └── README.md /.gitattributes: -------------------------------------------------------------------------------- 1 | neuralnet_pytorch/_version.py export-subst 2 | -------------------------------------------------------------------------------- /neuralnet_pytorch/zoo/__init__.py: -------------------------------------------------------------------------------- 1 | from .resnet import * 2 | from .vgg import * 3 | -------------------------------------------------------------------------------- /docs/source/_static/nnt-logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/justanhduc/neuralnet-pytorch/HEAD/docs/source/_static/nnt-logo.png -------------------------------------------------------------------------------- /neuralnet_pytorch/optim/lr_scheduler/__init__.py: -------------------------------------------------------------------------------- 1 | from .inverse_lr import InverseLR 2 | from .warm_restart import WarmRestart 3 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/pc2vox.py: -------------------------------------------------------------------------------- 1 | import neuralnet_pytorch.ext as ext 2 | 3 | __all__ = ['pc2vox'] 4 | 5 | pc2vox = ext.pc2vox_forward 6 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Created by .ignore support plugin (hsz.mobi) 2 | build 3 | dist 4 | examples/data 5 | examples/results 6 | neuralnet_pytorch.egg-info 7 | -------------------------------------------------------------------------------- /tests/test_files/test.gin: -------------------------------------------------------------------------------- 1 | net.dtype = %float32 2 | net.activation = @tanh 3 | net.loss = @l1loss 4 | net.optimizer = @adamax 5 | net.scheduler = @exp_lr 6 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include *.md 2 | include LICENSE.txt 3 | include versioneer.py 4 | graft neuralnet_pytorch/extensions/include 5 | include neuralnet_pytorch/_version.py 6 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/bpd.py: -------------------------------------------------------------------------------- 1 | import neuralnet_pytorch.ext as ext 2 | 3 | __all__ = ['batch_pairwise_dist'] 4 | 5 | batch_pairwise_dist = ext.bpd_forward 6 | -------------------------------------------------------------------------------- /neuralnet_pytorch/optim/__init__.py: -------------------------------------------------------------------------------- 1 | from .adabound import AdaBound 2 | from .nadam import NAdam 3 | from .lookahead import Lookahead 4 | from . import lr_scheduler 5 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/include/bpd.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | torch::Tensor batch_pairwise_distance_forward(torch::Tensor x, torch::Tensor y); 5 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch 2 | torchvision 3 | numpy 4 | scipy 5 | visdom 6 | matplotlib <= 3.1.0 7 | slackclient 8 | tensorboard 9 | imageio 10 | gin-config 11 | future 12 | sympy 13 | git-python -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/__init__.py: -------------------------------------------------------------------------------- 1 | from .dist_chamfer import chamfer_distance 2 | from .dist_emd import earth_mover_distance 3 | from .bpd import batch_pairwise_dist 4 | from .pc2vox import pc2vox 5 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/include/pc2vox.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | torch::Tensor pointcloud_to_voxel_forward(torch::Tensor pc, int voxel_size, 5 | float grid_size, bool filter_outlier); 6 | -------------------------------------------------------------------------------- /neuralnet_pytorch/layers/__init__.py: -------------------------------------------------------------------------------- 1 | from .abstract import * 2 | from .adain import * 3 | from .aggregation import * 4 | from .blocks import * 5 | from .convolution import * 6 | from .normalization import * 7 | from .points import * 8 | from .resizing import * 9 | -------------------------------------------------------------------------------- /docs/source/manual/metrics.rst: -------------------------------------------------------------------------------- 1 | .. _metrics: 2 | .. currentmodule:: neuralnet_pytorch 3 | 4 | ================================ 5 | :mod:`metrics` -- Common Metrics 6 | ================================ 7 | 8 | .. automodule:: neuralnet_pytorch.metrics 9 | :members: 10 | -------------------------------------------------------------------------------- /docs/source/manual/monitor.rst: -------------------------------------------------------------------------------- 1 | .. _monitor: 2 | .. currentmodule:: neuralnet_pytorch 3 | 4 | ===================================================== 5 | :mod:`monitor` -- Monitor Training of Neural Networks 6 | ===================================================== 7 | 8 | .. autoclass:: neuralnet_pytorch.monitor.Monitor 9 | :members: 10 | -------------------------------------------------------------------------------- /neuralnet_pytorch/gin_nnt/__init__.py: -------------------------------------------------------------------------------- 1 | """Init file for Pytorch-specific Gin-Config package. Adapted from Gin-config""" 2 | 3 | try: 4 | from gin import * 5 | from neuralnet_pytorch.gin_nnt import external_configurables 6 | except ImportError: 7 | print('Please install Gin-config first via \'pip install gin-config\'') 8 | raise 9 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | license_file = LICENSE.txt 3 | 4 | [versioneer] 5 | vcs = git 6 | style = pep440 7 | versionfile_source = neuralnet_pytorch/_version.py 8 | versionfile_build = neuralnet_pytorch/_version.py 9 | tag_prefix = rel- 10 | 11 | [egg_info] 12 | tag_build = 13 | tag_date = 0 14 | 15 | [flake8] 16 | exclude = *__init__* 17 | max-line-length = 120 18 | -------------------------------------------------------------------------------- /docs/source/manual/track.rst: -------------------------------------------------------------------------------- 1 | .. _track: 2 | .. currentmodule:: neuralnet_pytorch 3 | 4 | ============================================ 5 | :mod:`monitor` -- Track Intermediate Results 6 | ============================================ 7 | 8 | .. autofunction:: neuralnet_pytorch.monitor.track 9 | .. autofunction:: neuralnet_pytorch.monitor.get_tracked_variables 10 | .. autofunction:: neuralnet_pytorch.monitor.eval_tracked_variables 11 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/include/chamfer_cuda.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | 5 | std::vector chamfer_forward(at::Tensor xyz1, at::Tensor xyz2); 6 | std::vector chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, 7 | at::Tensor graddist1, 8 | at::Tensor graddist2, 9 | at::Tensor idx1, at::Tensor idx2); 10 | -------------------------------------------------------------------------------- /.travis/travis_install.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | if test -e $HOME/miniconda2/envs/pyenv; then 3 | echo "pyenv already exists." 4 | else 5 | echo "Creating pyenv." 6 | if [[ $TRAVIS_PYTHON_VERSION == '3.6' ]]; then conda create --yes -q -n pyenv python=3.6 ; fi 7 | if [[ $TRAVIS_PYTHON_VERSION == '3.7' ]]; then conda create --yes -q -n pyenv python=3.7 ; fi 8 | if [[ $TRAVIS_PYTHON_VERSION == '3.8' ]]; then conda create --yes -q -n pyenv python=3.8 ; fi 9 | fi 10 | 11 | source activate pyenv 12 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/csrc/bpd.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include "bpd.h" 4 | 5 | torch::Tensor 6 | batch_pairwise_distance_forward(torch::Tensor x, torch::Tensor y) { 7 | auto xx = at::sum(at::pow(x, 2), -1); 8 | auto yy = at::sum(at::pow(y, 2), -1); 9 | auto xy = at::matmul(x, at::transpose(y, -1, -2)); 10 | 11 | auto rx = xx.unsqueeze(-2).expand_as(at::transpose(xy, -1, -2)); 12 | auto ry = yy.unsqueeze(-2).expand_as(xy); 13 | auto P = at::transpose(rx, -1, -2) + ry - 2. * xy; 14 | return P; 15 | } 16 | -------------------------------------------------------------------------------- /docs/source/help.rst: -------------------------------------------------------------------------------- 1 | Contact 2 | ======= 3 | 4 | I am `Nguyen Anh Duc`_, a PhD candidate in Computer Vision at Yonsei University. 5 | This package is written during my little free time when doing PhD. 6 | Most but not all the written modules are properly checked. 7 | Therefore, buggy performance is inevitable. 8 | No replacements or refunds are guaranteed for such situations. 9 | If you want to collaborate or contribute, feel free to make a PR/issue 10 | on Github or email me at adnguyen@yonsei.ac.kr. 11 | 12 | .. _Nguyen Anh Duc: https://justanhduc.github.io/ 13 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/csrc/utils.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include "utils.h" 4 | 5 | torch::Tensor 6 | ravel_index(at::Tensor indices, at::Tensor shape) 7 | { 8 | torch::Tensor linear = torch::zeros({ indices.size(0), 1 }) 9 | .to(indices.device().type()) 10 | .to(torch::kLong); 11 | for (int i = 0; i < indices.size(1); ++i) 12 | linear = linear + 13 | indices.slice(1, i, i + 1) * 14 | at::prod(shape.slice(0, i + 1, shape.size(0))); 15 | 16 | return linear.flatten(); 17 | } 18 | -------------------------------------------------------------------------------- /.travis/travis_before_install.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | # Install miniconda to avoid compiling scipy 3 | if test -e $HOME/miniconda2/bin ; then 4 | echo "miniconda already installed." 5 | else 6 | echo "Installing miniconda." 7 | rm -rf $HOME/miniconda2 8 | mkdir -p $HOME/download 9 | if [[ -d $HOME/download/miniconda.sh ]] ; then rm -rf $HOME/download/miniconda.sh ; fi 10 | wget -c https://repo.continuum.io/miniconda/Miniconda2-4.5.11-Linux-x86_64.sh -O $HOME/download/miniconda.sh 11 | chmod +x $HOME/download/miniconda.sh 12 | $HOME/download/miniconda.sh -b 13 | fi 14 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/include/emd_cuda.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | 5 | at::Tensor approx_match_forward(const at::Tensor xyz1, const at::Tensor xyz2); 6 | at::Tensor match_cost_forward(const at::Tensor xyz1, const at::Tensor xyz2, 7 | const at::Tensor match); 8 | std::vector match_cost_backward(const at::Tensor grad_cost, 9 | const at::Tensor xyz1, 10 | const at::Tensor xyz2, 11 | const at::Tensor match); 12 | -------------------------------------------------------------------------------- /docs/source/manual/optimization.rst: -------------------------------------------------------------------------------- 1 | .. _opt: 2 | .. currentmodule:: neuralnet_pytorch 3 | 4 | ================================================= 5 | :mod:`optimization` -- Extra Optimization Schemes 6 | ================================================= 7 | 8 | Extra Optimizers 9 | ================ 10 | 11 | .. autoclass:: neuralnet_pytorch.optim.AdaBound 12 | .. autoclass:: neuralnet_pytorch.optim.Lookahead 13 | .. autoclass:: neuralnet_pytorch.optim.NAdam 14 | 15 | Extra LR Schedulers 16 | =================== 17 | 18 | .. autoclass:: neuralnet_pytorch.optim.lr_scheduler.InverseLR 19 | .. autoclass:: neuralnet_pytorch.optim.lr_scheduler.WarmRestart 20 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line. 5 | SPHINXOPTS = 6 | SPHINXBUILD = sphinx-build 7 | SOURCEDIR = source 8 | BUILDDIR = build 9 | 10 | # Put it first so that "make" without argument is like "make help". 11 | help: 12 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 13 | 14 | .PHONY: help Makefile 15 | 16 | # Catch-all target: route all unknown targets to Sphinx using the new 17 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 18 | %: Makefile 19 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/include/utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | 5 | #define CHECK_CUDA(x) \ 6 | TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 7 | #define CHECK_CONTIGUOUS(x) \ 8 | TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 9 | #define CHECK_INPUT(x) \ 10 | CHECK_CUDA(x); \ 11 | CHECK_CONTIGUOUS(x) 12 | 13 | torch::Tensor ravel_index(at::Tensor indices, at::Tensor shape); 14 | -------------------------------------------------------------------------------- /neuralnet_pytorch/optim/lr_scheduler/inverse_lr.py: -------------------------------------------------------------------------------- 1 | import torch.optim as optim 2 | 3 | 4 | class InverseLR(optim.lr_scheduler.LambdaLR): 5 | """Decreases lr every iteration by the inverse of gamma times iteration plus 1. 6 | :math:`\\text{lr} = \\text{lr} / (1 + \\gamma * t)`. 7 | 8 | Parameters 9 | ---------- 10 | optimizer 11 | wrapped optimizer. 12 | gamma 13 | decrease coefficient. 14 | last_epoch : int 15 | the index of last epoch. Default: -1. 16 | """ 17 | 18 | def __init__(self, optimizer, gamma, last_epoch=-1): 19 | self.gamma = gamma 20 | super().__init__(optimizer, lambda it: 1. / (1. + gamma * it), last_epoch) 21 | -------------------------------------------------------------------------------- /neuralnet_pytorch/utils/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | log_formatter = logging.Formatter('%(asctime)s [%(threadName)-12.12s] [%(levelname)-5.5s] %(message)s') 4 | root_logger = logging.getLogger() 5 | root_logger.setLevel(logging.INFO) 6 | 7 | from .tensor_utils import * 8 | from .layer_utils import * 9 | from .numpy_utils import * 10 | from .activation_utils import * 11 | from .cv_utils import * 12 | from .misc_utils import * 13 | from .data_utils import * 14 | from .layer_utils import _make_input_shape 15 | 16 | import torch 17 | 18 | cuda_available = torch.cuda.is_available() 19 | _image_shape = _make_input_shape(1, 2) 20 | _matrix_shape = _make_input_shape(1, 0) 21 | _pointset_shape = _make_input_shape(2, 0) 22 | 23 | del torch 24 | del _make_input_shape 25 | del logging 26 | -------------------------------------------------------------------------------- /docs/source/manual/zoo.rst: -------------------------------------------------------------------------------- 1 | .. _zoo: 2 | .. currentmodule:: neuralnet_pytorch 3 | 4 | 5 | ================================================ 6 | :mod:`zoo` -- Model Zoo - Pretrained Neural Nets 7 | ================================================ 8 | 9 | .. contents:: Contents 10 | :depth: 4 11 | 12 | Residual Neural Networks 13 | ======================== 14 | 15 | Adapted PyTorch's officially pretrained `ResNets `_ on ImageNet. 16 | 17 | .. automodule:: neuralnet_pytorch.zoo.resnet 18 | :members: 19 | :private-members: 20 | 21 | VGG 22 | === 23 | 24 | Adapted PyTorch's officially pretrained `VGG `_ on ImageNet. 25 | 26 | .. automodule:: neuralnet_pytorch.zoo.vgg 27 | :members: 28 | :private-members: 29 | -------------------------------------------------------------------------------- /neuralnet_pytorch/version.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, print_function, division 2 | 3 | from neuralnet_pytorch._version import get_versions 4 | 5 | author = 'DUC NGUYEN' 6 | FALLBACK_VERSION = "1.0.0+unknown" 7 | 8 | info = get_versions() 9 | if info['error'] is not None: 10 | info['version'] = FALLBACK_VERSION 11 | 12 | full_version = info['version'] 13 | git_revision = info['full-revisionid'] 14 | del get_versions 15 | 16 | short_version = full_version.split('+')[0] 17 | 18 | 19 | # This tries to catch a tag like beta2, rc1, ... 20 | try: 21 | int(short_version.split('.')[0]) 22 | release = True 23 | except ValueError: 24 | release = False 25 | 26 | if release and info['error'] is None: 27 | version = short_version 28 | else: 29 | version = full_version 30 | del info 31 | -------------------------------------------------------------------------------- /examples/configs/cifar.gin: -------------------------------------------------------------------------------- 1 | Classifier.name = 'resnet' 2 | Classifier.model = @resnet101 3 | Classifier.dataset = 'cifar10' 4 | 5 | Classifier.optimizer = @sgd # fused* works only on gpu0 due to NVIDIA's bug https://github.com/NVIDIA/apex/issues/319 6 | Classifier.scheduler = @multistep_lr 7 | Classifier.lr = 1e-3 8 | Classifier.weight_decay = 5e-4 9 | Classifier.momentum = 0.9 10 | Classifier.gamma = .1 11 | Classifier.milestones = (50, 75) 12 | 13 | Classifier.bs = 32 14 | Classifier.n_epochs = 100 15 | Classifier.start_epoch = None 16 | Classifier.print_freq = 500 17 | Classifier.val_freq = 500 18 | 19 | Classifier.use_jit = False 20 | Classifier.use_amp = False # works only on gpu0 due to NVIDIA's bug https://github.com/NVIDIA/apex/issues/319 21 | Classifier.opt_level='O1' 22 | 23 | Classifier.checkpoint_folder = None 24 | Classifier.version = -1 -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=source 11 | set BUILDDIR=build 12 | 13 | if "%1" == "" goto help 14 | 15 | %SPHINXBUILD% >NUL 2>NUL 16 | if errorlevel 9009 ( 17 | echo. 18 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 19 | echo.installed, then set the SPHINXBUILD environment variable to point 20 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 21 | echo.may add the Sphinx directory to PATH. 22 | echo. 23 | echo.If you don't have Sphinx installed, grab it from 24 | echo.http://sphinx-doc.org/ 25 | exit /b 1 26 | ) 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/csrc/bindings.cpp: -------------------------------------------------------------------------------- 1 | #include "bpd.h" 2 | #include "chamfer_cuda.h" 3 | #include "emd_cuda.h" 4 | #include "pc2vox.h" 5 | #include "utils.h" 6 | 7 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 8 | { 9 | m.def("chamfer_forward", &chamfer_forward, "chamfer forward (CUDA)"); 10 | m.def("chamfer_backward", &chamfer_backward, "chamfer backward (CUDA)"); 11 | m.def("approx_match_forward", &approx_match_forward, 12 | "ApproxMatch forward (CUDA)"); 13 | m.def("match_cost_forward", &match_cost_forward, "MatchCost forward (CUDA)"); 14 | m.def("match_cost_backward", &match_cost_backward, 15 | "MatchCost backward (CUDA)"); 16 | m.def("bpd_forward", &batch_pairwise_distance_forward, 17 | "batch pairwise distance forward"); 18 | m.def("pc2vox_forward", &pointcloud_to_voxel_forward, 19 | "pointcloud to voxel forward"); 20 | m.def("ravel_index_forward", &ravel_index, "ravel index forward"); 21 | } 22 | -------------------------------------------------------------------------------- /docs/source/manual/resizing.rst: -------------------------------------------------------------------------------- 1 | .. _resizing: 2 | .. currentmodule:: neuralnet_pytorch 3 | 4 | ================================== 5 | :mod:`resizing` -- Resizing Layers 6 | ================================== 7 | 8 | .. contents:: Contents 9 | :depth: 4 10 | 11 | Extended Resizing Layers 12 | ======================== 13 | 14 | .. autoclass:: neuralnet_pytorch.resizing.Interpolate 15 | .. autoclass:: neuralnet_pytorch.resizing.AvgPool2d 16 | .. autoclass:: neuralnet_pytorch.resizing.MaxPool2d 17 | 18 | Custom Resizing Layers 19 | ====================== 20 | 21 | .. autoclass:: neuralnet_pytorch.resizing.GlobalAvgPool2D 22 | .. autoclass:: neuralnet_pytorch.resizing.Cat 23 | .. autoclass:: neuralnet_pytorch.resizing.ConcurrentCat 24 | .. autoclass:: neuralnet_pytorch.resizing.SequentialCat 25 | .. autoclass:: neuralnet_pytorch.resizing.Reshape 26 | .. autoclass:: neuralnet_pytorch.resizing.Flatten 27 | .. autoclass:: neuralnet_pytorch.resizing.DimShuffle 28 | -------------------------------------------------------------------------------- /docs/source/manual/normalization.rst: -------------------------------------------------------------------------------- 1 | .. _normalization: 2 | .. currentmodule:: neuralnet_pytorch 3 | 4 | ============================================ 5 | :mod:`normalization` -- Normalization Layers 6 | ============================================ 7 | 8 | .. contents:: Contents 9 | :depth: 4 10 | 11 | Extended Normalization Layers 12 | ============================= 13 | 14 | .. autoclass:: neuralnet_pytorch.layers.BatchNorm1d 15 | .. autoclass:: neuralnet_pytorch.layers.BatchNorm2d 16 | .. autoclass:: neuralnet_pytorch.layers.LayerNorm 17 | .. autoclass:: neuralnet_pytorch.layers.InstanceNorm1d 18 | .. autoclass:: neuralnet_pytorch.layers.InstanceNorm2d 19 | .. autoclass:: neuralnet_pytorch.layers.GroupNorm 20 | 21 | Custom Lormalization Layers 22 | =========================== 23 | 24 | .. autoclass:: neuralnet_pytorch.layers.FeatureNorm1d 25 | .. autoclass:: neuralnet_pytorch.layers.AdaIN 26 | .. autoclass:: neuralnet_pytorch.layers.MultiModuleAdaIN 27 | .. autoclass:: neuralnet_pytorch.layers.MultiInputAdaIN 28 | -------------------------------------------------------------------------------- /neuralnet_pytorch/optim/lr_scheduler/warm_restart.py: -------------------------------------------------------------------------------- 1 | from torch.optim.lr_scheduler import CosineAnnealingLR 2 | 3 | 4 | class WarmRestart(CosineAnnealingLR): 5 | """ 6 | Step should be used in batch iteration loop. 7 | Putting step in the epoch loop results in wrong behavior of the restart. 8 | One must not pass the iteration number to step. 9 | 10 | Parameters 11 | ---------- 12 | optimizer 13 | wrapped optimizer. 14 | T_max 15 | maximum number of iterations. 16 | T_mul 17 | multiplier for `T_max`. 18 | eta_min 19 | minimum learning rate. 20 | Default: 0. 21 | last_epoch 22 | the index of last epoch. 23 | Default: -1. 24 | """ 25 | 26 | def __init__(self, optimizer, T_max, T_mul=1, eta_min=0, last_epoch=-1): 27 | self.T_mul = T_mul 28 | super().__init__(optimizer, T_max, eta_min, last_epoch) 29 | 30 | def get_lr(self): 31 | if self.last_epoch > self.T_max: 32 | self.last_epoch = 0 33 | self.T_max *= self.T_mul 34 | return super().get_lr() 35 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Duc Nguyen 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/csrc/chamfer_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include "chamfer_cuda.h" 2 | 3 | #include "utils.h" 4 | 5 | std::vector chamfer_cuda_forward(at::Tensor xyz1, 6 | at::Tensor xyz2); 7 | std::vector chamfer_cuda_backward( 8 | at::Tensor xyz1, at::Tensor xyz2, at::Tensor graddist1, at::Tensor graddist2, 9 | at::Tensor idx1, at::Tensor idx2); 10 | 11 | std::vector 12 | chamfer_forward(at::Tensor xyz1, at::Tensor xyz2) 13 | { 14 | CHECK_EQ(xyz1.size(0), xyz2.size(0)); 15 | CHECK_EQ(xyz1.size(2), xyz2.size(2)); 16 | CHECK_INPUT(xyz1); 17 | CHECK_INPUT(xyz2); 18 | return chamfer_cuda_forward(xyz1, xyz2); 19 | } 20 | 21 | std::vector 22 | chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor graddist1, 23 | at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) 24 | { 25 | CHECK_EQ(xyz1.size(0), xyz2.size(0)); 26 | CHECK_EQ(xyz1.size(1), idx1.size(1)); 27 | CHECK_EQ(xyz2.size(1), idx2.size(1)); 28 | CHECK_INPUT(xyz1); 29 | CHECK_INPUT(xyz2); 30 | return chamfer_cuda_backward(xyz1, xyz2, graddist1, graddist2, idx1, idx2); 31 | } 32 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/dist_emd.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | import neuralnet_pytorch.ext as ext 3 | 4 | __all__ = ['earth_mover_distance'] 5 | 6 | 7 | class EarthMoverDistanceFunction(T.autograd.Function): 8 | @staticmethod 9 | def forward(ctx, xyz1, xyz2): 10 | xyz1 = xyz1.contiguous() 11 | xyz2 = xyz2.contiguous() 12 | assert xyz1.is_cuda and xyz2.is_cuda, "Only support cuda currently." 13 | match = ext.approx_match_forward(xyz1, xyz2) 14 | cost = ext.match_cost_forward(xyz1, xyz2, match) 15 | ctx.save_for_backward(xyz1, xyz2, match) 16 | return cost 17 | 18 | @staticmethod 19 | def backward(ctx, grad_cost): 20 | xyz1, xyz2, match = ctx.saved_tensors 21 | grad_cost = grad_cost.contiguous() 22 | grad_xyz1, grad_xyz2 = ext.match_cost_backward(grad_cost, xyz1, xyz2, match) 23 | return grad_xyz1, grad_xyz2 24 | 25 | 26 | def earth_mover_distance(xyz1, xyz2): 27 | if xyz1.dim() == 2: 28 | xyz1 = xyz1.unsqueeze(0) 29 | 30 | if xyz2.dim() == 2: 31 | xyz2 = xyz2.unsqueeze(0) 32 | 33 | cost = EarthMoverDistanceFunction.apply(xyz1, xyz2) 34 | return cost 35 | -------------------------------------------------------------------------------- /docs/source/license.rst: -------------------------------------------------------------------------------- 1 | License 2 | ======= 3 | 4 | MIT License 5 | 6 | Copyright (c) 2019 Duc Nguyen 7 | 8 | Permission is hereby granted, free of charge, to any person obtaining a copy 9 | of this software and associated documentation files (the "Software"), to deal 10 | in the Software without restriction, including without limitation the rights 11 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 12 | copies of the Software, and to permit persons to whom the Software is 13 | furnished to do so, subject to the following conditions: 14 | 15 | The above copyright notice and this permission notice shall be included in all 16 | copies or substantial portions of the Software. 17 | 18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 24 | SOFTWARE. 25 | -------------------------------------------------------------------------------- /docs/source/manual/index.rst: -------------------------------------------------------------------------------- 1 | .. _manual: 2 | 3 | ================ 4 | Reference Manual 5 | ================ 6 | 7 | To check if you have installed Neuralnet-pytorch correctly, try 8 | 9 | >>> import neuralnet_pytorch as nnt 10 | 11 | If there is any error, please check your installation again (see :ref:`installation`). 12 | 13 | The basic contents of our package are as follows. 14 | 15 | Neuralnet Layers 16 | ---------------- 17 | .. toctree:: 18 | :maxdepth: 4 19 | 20 | layers 21 | normalization 22 | resizing 23 | 24 | Monitor Network Training 25 | ------------------------ 26 | .. toctree:: 27 | :maxdepth: 4 28 | 29 | monitor 30 | track 31 | 32 | Optimization 33 | ------------ 34 | .. toctree:: 35 | :maxdepth: 4 36 | 37 | optimization 38 | 39 | Additional Metrics 40 | ------------------ 41 | .. toctree:: 42 | :maxdepth: 4 43 | 44 | metrics 45 | 46 | Extras 47 | ------ 48 | .. toctree:: 49 | :maxdepth: 4 50 | 51 | utilities 52 | 53 | Gin-config 54 | ---------- 55 | .. toctree:: 56 | :maxdepth: 4 57 | 58 | gin 59 | 60 | Pretrained Networks 61 | ------------------- 62 | .. toctree:: 63 | :maxdepth: 4 64 | 65 | zoo 66 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/dist_chamfer.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from torch.autograd import Function 3 | import neuralnet_pytorch.ext as ext 4 | 5 | __all__ = ['chamfer_distance'] 6 | 7 | 8 | class ChamferFunction(Function): 9 | """ 10 | Chamfer's distance module @thibaultgroueix 11 | GPU tensors only 12 | """ 13 | 14 | @staticmethod 15 | def forward(ctx, xyz1, xyz2): 16 | batchsize, n, _ = xyz1.size() 17 | _, m, _ = xyz2.size() 18 | 19 | dist1, dist2, idx1, idx2 = ext.chamfer_forward(xyz1, xyz2) 20 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 21 | return dist1, dist2 22 | 23 | @staticmethod 24 | def backward(ctx, graddist1, graddist2): 25 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 26 | graddist1 = graddist1.contiguous() 27 | graddist2 = graddist2.contiguous() 28 | grad_xyz1, grad_xyz2 = ext.chamfer_backward(xyz1, xyz2, graddist1, graddist2, idx1, idx2) 29 | return grad_xyz1, grad_xyz2 30 | 31 | 32 | class ChamferDistance(nn.Module): 33 | def __init__(self): 34 | super(ChamferDistance, self).__init__() 35 | 36 | def forward(self, input1, input2): 37 | return ChamferFunction.apply(input1, input2) 38 | 39 | 40 | chamfer_distance = ChamferDistance() 41 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | cache: 2 | directories: 3 | - $HOME/.cache/pip 4 | - $HOME/download # Sufficient to add miniconda.sh to TRAVIS cache. 5 | - $HOME/miniconda2 # Add the installation to TRAVIS cache. 6 | 7 | dist: xenial 8 | 9 | language: python 10 | 11 | python: 12 | - "3.6" 13 | - "3.7" 14 | - "3.8" 15 | 16 | before_install: 17 | - ./.travis/travis_before_install.sh 18 | - export PATH=/home/travis/miniconda2/bin:$PATH 19 | 20 | install: 21 | - ./.travis/travis_install.sh 22 | - source activate pyenv 23 | - conda install -y pytorch torchvision cpuonly -c pytorch 24 | - pip install pytest-cov pytest 25 | - pip install codecov 26 | - pip install flake8 27 | - pip install -r requirements.txt 28 | - pip install . 29 | 30 | jobs: 31 | - &normaltest 32 | stage: test 33 | python: "3.6" 34 | - <<: *normaltest 35 | python: "3.7" 36 | - <<: *normaltest 37 | python: "3.8" 38 | 39 | script: 40 | - python --version 41 | - uname -a 42 | - free -m 43 | - df -h 44 | - ulimit -a 45 | - echo "$PART" 46 | - python -c 'import torch; print(torch.__version__)' 47 | - python -c 'import neuralnet_pytorch; print(neuralnet_pytorch.__version__)' 48 | - python -m pytest -v --cov=neuralnet_pytorch/ 49 | - flake8 neuralnet_pytorch 50 | 51 | env: 52 | - CODECOV_TOKEN="41bdb269-8e93-42a6-a87d-65339e8145e9" 53 | 54 | after_failure: 55 | - cat /home/travis/.pip/pip.log 56 | 57 | after_success: 58 | - bash <(curl -s https://codecov.io/bash) -------------------------------------------------------------------------------- /docs/source/manual/utilities.rst: -------------------------------------------------------------------------------- 1 | .. _utilities: 2 | .. currentmodule:: neuralnet_pytorch 3 | 4 | ================================================= 5 | :mod:`utilities` -- Utility Functions and Classes 6 | ================================================= 7 | 8 | .. contents:: Contents 9 | :depth: 4 10 | 11 | Utilities for tensor manipulation 12 | ================================= 13 | 14 | .. automodule:: neuralnet_pytorch.utils.tensor_utils 15 | :members: 16 | :private-members: 17 | 18 | Utilities for Pytorch layers 19 | ============================ 20 | 21 | .. automodule:: neuralnet_pytorch.utils.layer_utils 22 | :members: 23 | :private-members: 24 | 25 | Utilities for data handling 26 | =========================== 27 | 28 | .. automodule:: neuralnet_pytorch.utils.data_utils 29 | :members: 30 | :private-members: 31 | 32 | Utilities for activation functions 33 | ================================== 34 | 35 | .. automodule:: neuralnet_pytorch.utils.activation_utils 36 | :members: 37 | :private-members: 38 | 39 | Utilities for computer vision applications 40 | ========================================== 41 | 42 | .. automodule:: neuralnet_pytorch.utils.cv_utils 43 | :members: 44 | :private-members: 45 | 46 | Utilities for Numpy arrays 47 | ========================== 48 | 49 | .. automodule:: neuralnet_pytorch.utils.numpy_utils 50 | :members: 51 | :private-members: 52 | 53 | Miscellaneous utilities 54 | ======================= 55 | 56 | .. automodule:: neuralnet_pytorch.utils.misc_utils 57 | :members: 58 | :private-members: 59 | -------------------------------------------------------------------------------- /neuralnet_pytorch/utils/misc_utils.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | 3 | from . import root_logger 4 | 5 | __all__ = ['time_cuda_module', 'slack_message'] 6 | 7 | 8 | def time_cuda_module(f, *args, **kwargs): 9 | """ 10 | Measures the time taken by a Pytorch module. 11 | 12 | :param f: 13 | a Pytorch module. 14 | :param args: 15 | arguments to be passed to `f`. 16 | :param kwargs: 17 | keyword arguments to be passed to `f`. 18 | :return: 19 | the time (in second) that `f` takes. 20 | """ 21 | 22 | start = T.cuda.Event(enable_timing=True) 23 | end = T.cuda.Event(enable_timing=True) 24 | 25 | start.record() 26 | f(*args, **kwargs) 27 | end.record() 28 | 29 | # Waits for everything to finish running 30 | T.cuda.synchronize() 31 | 32 | total = start.elapsed_time(end) 33 | root_logger.info('Took %fms' % total) 34 | return total 35 | 36 | 37 | def slack_message(username: str, message: str, channel: str, token: str, **kwargs): 38 | """ 39 | Sends a slack message to the specified chatroom. 40 | 41 | :param username: 42 | Slack username. 43 | :param message: 44 | message to be sent. 45 | :param channel: 46 | Slack channel. 47 | :param token: 48 | Slack chatroom token. 49 | :param kwargs: 50 | additional keyword arguments to slack's :meth:`api_call`. 51 | :return: 52 | ``None``. 53 | """ 54 | 55 | try: 56 | from slackclient import SlackClient 57 | except (ModuleNotFoundError, ImportError): 58 | from slack import RTMClient as SlackClient 59 | 60 | sc = SlackClient(token) 61 | sc.api_call('chat.postMessage', channel=channel, text=message, username=username, **kwargs) 62 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/csrc/emd_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include "emd_cuda.h" 2 | #include "utils.h" 3 | 4 | at::Tensor approx_match_cuda_forward(const at::Tensor xyz1, 5 | const at::Tensor xyz2); 6 | at::Tensor match_cost_cuda_forward(const at::Tensor xyz1, const at::Tensor xyz2, 7 | const at::Tensor match); 8 | std::vector match_cost_cuda_backward(const at::Tensor grad_cost, 9 | const at::Tensor xyz1, 10 | const at::Tensor xyz2, 11 | const at::Tensor match); 12 | 13 | at::Tensor 14 | approx_match_forward(const at::Tensor xyz1, const at::Tensor xyz2) 15 | { 16 | CHECK_EQ(xyz1.size(0), xyz2.size(0)); 17 | CHECK_EQ(xyz1.size(2), 3); 18 | CHECK_EQ(xyz2.size(2), 3); 19 | CHECK_INPUT(xyz1); 20 | CHECK_INPUT(xyz2); 21 | return approx_match_cuda_forward(xyz1, xyz2); 22 | } 23 | 24 | at::Tensor 25 | match_cost_forward(const at::Tensor xyz1, const at::Tensor xyz2, 26 | const at::Tensor match) 27 | { 28 | CHECK_EQ(xyz1.size(0), xyz2.size(0)); 29 | CHECK_EQ(xyz1.size(2), 3); 30 | CHECK_EQ(xyz2.size(2), 3); 31 | CHECK_INPUT(xyz1); 32 | CHECK_INPUT(xyz2); 33 | return match_cost_cuda_forward(xyz1, xyz2, match); 34 | } 35 | 36 | std::vector 37 | match_cost_backward(const at::Tensor grad_cost, const at::Tensor xyz1, 38 | const at::Tensor xyz2, const at::Tensor match) 39 | { 40 | CHECK_EQ(xyz1.size(0), xyz2.size(0)); 41 | CHECK_EQ(xyz1.size(2), 3); 42 | CHECK_EQ(xyz2.size(2), 3); 43 | CHECK_INPUT(xyz1); 44 | CHECK_INPUT(xyz2); 45 | return match_cost_cuda_backward(grad_cost, xyz1, xyz2, match); 46 | } 47 | -------------------------------------------------------------------------------- /neuralnet_pytorch/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | 3 | minimum_required = '1.0.0' 4 | 5 | 6 | # Ensure Pytorch is importable and its version is sufficiently recent. This 7 | # needs to happen before anything else, since the imports below will try to 8 | # import Pytorch, too. 9 | def _ensure_pt_install(): # pylint: disable=g-statement-before-imports 10 | """Attempt to import Pytorch, and ensure its version is sufficient. 11 | Raises: 12 | ImportError: if either Pytorch is not importable or its version is 13 | inadequate. 14 | """ 15 | try: 16 | import torch 17 | except ImportError: 18 | # Print more informative error message, then reraise. 19 | print('\n\nFailed to import Pytorch. ' 20 | 'To use neuralnet-pytorch, please install ' 21 | 'Pytorch (> %s) by following instructions at ' 22 | 'https://pytorch.org/get-started/locally/.\n\n' % minimum_required) 23 | raise 24 | 25 | del torch 26 | 27 | 28 | _ensure_pt_install() 29 | 30 | # Cleanup symbols to avoid polluting namespace. 31 | del minimum_required 32 | import sys as _sys 33 | 34 | for symbol in ['_ensure_pt_install', '_sys']: 35 | delattr(_sys.modules[__name__], symbol) 36 | 37 | try: 38 | import neuralnet_pytorch.ext as ext 39 | cuda_ext_available = True 40 | del ext 41 | except ModuleNotFoundError: 42 | cuda_ext_available = False 43 | 44 | from . import utils 45 | from .utils import DataLoader, DataPrefetcher, cuda_available, function 46 | from .layers import * 47 | from .metrics import * 48 | from .monitor import * 49 | from . import optim 50 | 51 | from .version import author as __author__ 52 | from ._version import get_versions 53 | 54 | __version__ = get_versions()['version'] 55 | del get_versions 56 | -------------------------------------------------------------------------------- /docs/source/index.rst: -------------------------------------------------------------------------------- 1 | Welcome to Neuralnet-Pytorch's documentation! 2 | ============================================= 3 | 4 | Personally, going from Theano to Pytorch is pretty much like 5 | time traveling from 90s to the modern day. 6 | However, we feel like despite having a lot of bells and whistles, 7 | Pytorch is still missing many elements 8 | that are confirmed to never be added to the library. 9 | Therefore, this library is written to supplement more features 10 | to the current magical Pytorch. 11 | All the modules in the package directly subclass 12 | the corresponding modules from Pytorch, 13 | so everything should still be familiar. 14 | For example, the following snippet in Pytorch :: 15 | 16 | from torch import nn 17 | model = nn.Sequential( 18 | nn.Conv2d(1, 20, 5, padding=2), 19 | nn.ReLU(), 20 | nn.Conv2d(20, 64, 5, padding=2), 21 | nn.ReLU() 22 | ) 23 | 24 | 25 | can be rewritten in Neuralnet-pytorch as :: 26 | 27 | import neuralnet_pytorch as nnt 28 | model = nnt.Sequential( 29 | nnt.Conv2d(1, 20, 5, padding='half', activation='relu'), 30 | nnt.Conv2d(20, 64, 5, padding='half', activation='relu') 31 | ) 32 | 33 | which is the same as the native Pytorch, or :: 34 | 35 | import neuralnet_pytorch as nnt 36 | model = nnt.Sequential(input_shape=1) 37 | model.add_module('conv1', nnt.Conv2d(model.output_shape, 20, 5, padding='half', activation='relu')) 38 | model.add_module('conv2', nnt.Conv2d(model.output_shape, 64, 5, padding='half', activation='relu')) 39 | 40 | which frees you from doing a lot of manual calculations when adding one layer on top of another. Theano folks will also 41 | find some reminiscence as many functions are highly inspired by Theano. 42 | 43 | 44 | .. toctree:: 45 | :maxdepth: 4 46 | :caption: Overview: 47 | 48 | installation 49 | manual/index 50 | 51 | .. toctree:: 52 | :maxdepth: 3 53 | :caption: Misc Notes: 54 | 55 | license 56 | help 57 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/csrc/pc2vox.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include "pc2vox.h" 4 | #include "utils.h" 5 | 6 | torch::Tensor 7 | pointcloud_to_voxel_forward(torch::Tensor pc, int voxel_size, float grid_size, 8 | bool filter_outlier) 9 | { 10 | auto b = pc.size(0), n = pc.size(1); 11 | float half_size = grid_size / 2.; 12 | auto valid = torch::ones({ b * n }).to(pc.device().type()).to(torch::kLong); 13 | 14 | auto pc_grid = (pc + half_size) * (voxel_size - 1.); 15 | auto indices_floor = at::floor(pc_grid); 16 | auto indices = indices_floor.to(torch::kLong); 17 | 18 | auto batch_indices = 19 | torch::arange(b).to(pc.device().type()).to(indices.dtype()); 20 | batch_indices = batch_indices.unsqueeze(1).unsqueeze(2); 21 | batch_indices = batch_indices.expand({ b, n, 1 }); 22 | indices = at::cat({ batch_indices, indices }, 2); 23 | indices = at::reshape(indices, { -1, 4 }); 24 | auto r = pc_grid - indices_floor; 25 | std::vector rr{ 1. - r, r }; 26 | 27 | if (filter_outlier) { 28 | valid = 29 | at::all(at::__and__(at::ge(pc, -half_size), at::le(pc, half_size)), 2); 30 | valid = valid.flatten(); 31 | indices = indices.index(valid); 32 | } 33 | 34 | std::vector output_shape{ b, voxel_size, voxel_size, voxel_size }; 35 | at::Tensor output_shape_tensor = 36 | torch::tensor(output_shape).to(pc.device().type()).to(torch::kLong); 37 | auto voxel = torch::zeros(output_shape, pc.type()).flatten(); 38 | for (int k = 0; k < 2; ++k) { 39 | for (int j = 0; j < 2; ++j) { 40 | for (int i = 0; i < 2; ++i) { 41 | auto updates = 42 | rr[k].slice(2, 0, 1) * rr[j].slice(2, 1, 2) * rr[i].slice(2, 2, 3); 43 | updates = updates.flatten(); 44 | if (filter_outlier) 45 | updates = updates.index(valid); 46 | 47 | std::vector shift{ 0, k, j, i }; 48 | at::Tensor indices_shift = torch::tensor(shift) 49 | .to(pc.device().type()) 50 | .to(torch::kLong) 51 | .unsqueeze(0); 52 | auto indices_tmp = indices + indices_shift; 53 | 54 | auto linear_indices = ravel_index(indices_tmp, output_shape_tensor); 55 | voxel = scatter_add(voxel, 0, linear_indices, updates); 56 | } 57 | } 58 | } 59 | voxel = voxel.reshape(output_shape); 60 | voxel = at::clamp(voxel, 0., 1.); 61 | return voxel; 62 | } 63 | -------------------------------------------------------------------------------- /neuralnet_pytorch/utils/activation_utils.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | import torch.nn.functional as F 3 | from functools import partial, update_wrapper 4 | 5 | __all__ = ['relu', 'linear', 'lrelu', 'tanh', 'sigmoid', 'elu', 'selu', 'softmax', 'function'] 6 | 7 | 8 | def relu(x: T.Tensor, **kwargs): 9 | """ 10 | ReLU activation. 11 | """ 12 | 13 | return T.relu(x) 14 | 15 | 16 | def linear(x: T.Tensor, **kwargs): 17 | """ 18 | Linear activation. 19 | """ 20 | 21 | return x 22 | 23 | 24 | def lrelu(x: T.Tensor, **kwargs): 25 | """ 26 | Leaky ReLU activation. 27 | """ 28 | 29 | return F.leaky_relu(x, kwargs.get('negative_slope', .2), kwargs.get('inplace', False)) 30 | 31 | 32 | def tanh(x: T.Tensor, **kwargs): 33 | """ 34 | Hyperbolic tangent activation. 35 | """ 36 | 37 | return T.tanh(x) 38 | 39 | 40 | def sigmoid(x: T.Tensor, **kwargs): 41 | """ 42 | Sigmoid activation. 43 | """ 44 | 45 | return T.sigmoid(x) 46 | 47 | 48 | def elu(x: T.Tensor, **kwargs): 49 | """ 50 | ELU activation. 51 | """ 52 | 53 | return F.elu(x, kwargs.get('alpha', 1.), kwargs.get('inplace', False)) 54 | 55 | 56 | def softmax(x: T.Tensor, **kwargs): 57 | """ 58 | Softmax activation. 59 | """ 60 | 61 | return T.softmax(x, kwargs.get('dim', None)) 62 | 63 | 64 | def selu(x: T.Tensor, **kwargs): 65 | """ 66 | SELU activation. 67 | """ 68 | 69 | return T.selu(x) 70 | 71 | 72 | act = { 73 | 'relu': relu, 74 | 'linear': linear, 75 | None: linear, 76 | 'lrelu': lrelu, 77 | 'tanh': tanh, 78 | 'sigmoid': sigmoid, 79 | 'elu': elu, 80 | 'softmax': softmax, 81 | 'selu': selu 82 | } 83 | 84 | 85 | def function(activation, **kwargs): 86 | """ 87 | returns the `activation`. Can be ``str`` or ``callable``. 88 | For ``str``, possible choices are 89 | ``None``, ``'linear'``, ``'relu'``, ``'lrelu'``, 90 | ``'tanh'``, ``'sigmoid'``, ``'elu'``, ``'softmax'``, 91 | and ``'selu'``. 92 | 93 | :param activation: 94 | name of the activation function. 95 | :return: 96 | activation function 97 | """ 98 | 99 | func = partial(activation, **kwargs) if callable(activation) else partial(act[activation], **kwargs) 100 | return update_wrapper(func, activation) if callable(activation) else update_wrapper(func, act[activation]) 101 | -------------------------------------------------------------------------------- /docs/source/manual/layers.rst: -------------------------------------------------------------------------------- 1 | .. _layers: 2 | .. currentmodule:: neuralnet_pytorch 3 | 4 | ============================= 5 | :mod:`layers` -- Basic Layers 6 | ============================= 7 | 8 | .. module:: neuralnet_pytorch.layers 9 | :platform: Unix, Windows 10 | :synopsis: Basics layers in Deep Neural Networks 11 | .. moduleauthor:: Duc Nguyen 12 | 13 | This section describes all the backbone modules of Neuralnet-pytorch. 14 | 15 | .. contents:: Contents 16 | :depth: 4 17 | 18 | Abstract Layers 19 | =============== 20 | 21 | Attributes 22 | ---------- 23 | The following classes equip the plain ``torch`` modules with more bells and whistles. 24 | Also, some features are deeply integrated into Neuralnet-pytorch, 25 | which enables faster and more convenient training and testing of your neural networks. 26 | 27 | .. autoclass:: neuralnet_pytorch.layers.abstract._LayerMethod 28 | :members: 29 | .. autoclass:: neuralnet_pytorch.layers.MultiSingleInputModule 30 | .. autoclass:: neuralnet_pytorch.layers.MultiMultiInputModule 31 | 32 | 33 | Extended Pytorch Abstract Layers 34 | -------------------------------- 35 | 36 | .. autoclass:: neuralnet_pytorch.layers.Module 37 | .. autoclass:: neuralnet_pytorch.layers.Sequential 38 | 39 | Quick-and-dirty Layers 40 | ---------------------- 41 | 42 | .. autodecorator:: neuralnet_pytorch.layers.wrapper 43 | .. autoclass:: neuralnet_pytorch.layers.Lambda 44 | 45 | Common Layers 46 | ============= 47 | 48 | Extended Pytorch Common Layers 49 | ------------------------------ 50 | 51 | .. autoclass:: neuralnet_pytorch.layers.Conv2d 52 | .. autoclass:: neuralnet_pytorch.layers.ConvTranspose2d 53 | .. autoclass:: neuralnet_pytorch.layers.FC 54 | .. autoclass:: neuralnet_pytorch.layers.Softmax 55 | 56 | Extra Layers 57 | ------------ 58 | 59 | .. autoclass:: neuralnet_pytorch.layers.Activation 60 | .. autoclass:: neuralnet_pytorch.layers.ConvNormAct 61 | .. autoclass:: neuralnet_pytorch.layers.DepthwiseSepConv2D 62 | .. autoclass:: neuralnet_pytorch.layers.FCNormAct 63 | .. autoclass:: neuralnet_pytorch.layers.ResNetBasicBlock 64 | .. autoclass:: neuralnet_pytorch.layers.ResNetBottleneckBlock 65 | .. autoclass:: neuralnet_pytorch.layers.StackingConv 66 | 67 | Graph Learning Layers 68 | --------------------- 69 | 70 | .. autoclass:: neuralnet_pytorch.layers.GraphConv 71 | .. autoclass:: neuralnet_pytorch.layers.BatchGraphConv 72 | .. autoclass:: neuralnet_pytorch.layers.GraphXConv 73 | 74 | Multi-module Layers 75 | ------------------- 76 | 77 | .. autoclass:: neuralnet_pytorch.layers.Sum 78 | .. autoclass:: neuralnet_pytorch.layers.ConcurrentSum 79 | .. autoclass:: neuralnet_pytorch.layers.SequentialSum 80 | -------------------------------------------------------------------------------- /docs/source/installation.rst: -------------------------------------------------------------------------------- 1 | Installation 2 | ============ 3 | 4 | .. contents:: 5 | :depth: 3 6 | :local: 7 | 8 | Requirements 9 | ------------ 10 | 11 | Pytorch 12 | ^^^^^^^ 13 | 14 | Neuralnet-pytorch is built on top of Pytorch, so obviously Pytorch is needed. 15 | Please refer to the official `Pytorch website `_ for installation details. 16 | 17 | 18 | Other dependencies 19 | ^^^^^^^^^^^^^^^^^^ 20 | 21 | In Neuralnet-pytorch, we use several backends to visualize training, so it is necessary to install 22 | some additional packages. For convenience, installing the package will install all the required 23 | dependencies. Optional dependencies can be installed as instructed below. 24 | 25 | 26 | Install Neuralnet-pytorch 27 | ------------------------- 28 | 29 | There are two ways to install Neuralnet-pytorch: via PyPi and Github. 30 | At the moment, the package is not available on Conda yet. 31 | 32 | From PyPi 33 | ^^^^^^^^^ 34 | 35 | The easiest and quickest way to get Neuralnet-pytorch is to install the package from Pypi. 36 | In a Terminal session, simply type :: 37 | 38 | pip install neuralnet-pytorch 39 | 40 | From Github 41 | ^^^^^^^^^^^ 42 | 43 | To install the bleeding-edge version, which is highly recommended, run :: 44 | 45 | pip install git+git://github.com/justanhduc/neuralnet-pytorch.git@master 46 | 47 | 48 | To install the package with optional dependencies, try :: 49 | 50 | pip install "neuralnet-pytorch[option] @ git+git://github.com/justanhduc/neuralnet-pytorch.git@master" 51 | 52 | in which ``option`` can be ``gin``/``geom``/``visdom``/``slack``. 53 | 54 | 55 | We also provide a version with some fancy Cuda/C++ implementations 56 | that are implemented or collected from various sources. To install this version, run :: 57 | 58 | pip install neuralnet-pytorch --cuda-ext 59 | 60 | Uninstall Neuralnet-pytorch 61 | --------------------------- 62 | 63 | Simply use pip to uninstall the package :: 64 | 65 | pip uninstall neuralnet-pytorch 66 | 67 | Why would you want to do that anyway? 68 | 69 | Upgrade Neuralnet-pytorch 70 | ------------------------- 71 | 72 | Use pip with ``-U`` or ``--upgrade`` option :: 73 | 74 | pip install -U neuralnet-pytorch 75 | 76 | However, for maximal experience, please considering using the bleeding-edge version on Github. 77 | 78 | Reinstall Neuralnet-pytorch 79 | --------------------------- 80 | 81 | If you want to reinstall Neuralnet-pytorch, please uninstall and then install it again. 82 | When reinstalling the package, we recommend to use ``--no-cache-dir`` option as pip caches 83 | the previously built binaries :: 84 | 85 | pip uninstall neuralnet-pytorch 86 | pip install neuralnet-pytorch --no-cache-dir 87 | 88 | 89 | -------------------------------------------------------------------------------- /.clang-format: -------------------------------------------------------------------------------- 1 | --- 2 | Language: Cpp 3 | # BasedOnStyle: Mozilla 4 | AccessModifierOffset: -2 5 | AlignAfterOpenBracket: Align 6 | AlignConsecutiveAssignments: false 7 | AlignConsecutiveDeclarations: false 8 | AlignEscapedNewlinesLeft: false 9 | AlignOperands: true 10 | AlignTrailingComments: true 11 | AllowAllParametersOfDeclarationOnNextLine: false 12 | AllowShortBlocksOnASingleLine: false 13 | AllowShortCaseLabelsOnASingleLine: false 14 | AllowShortFunctionsOnASingleLine: Inline 15 | AllowShortIfStatementsOnASingleLine: false 16 | AllowShortLoopsOnASingleLine: false 17 | AlwaysBreakAfterDefinitionReturnType: TopLevel 18 | AlwaysBreakAfterReturnType: TopLevelDefinitions 19 | AlwaysBreakBeforeMultilineStrings: false 20 | AlwaysBreakTemplateDeclarations: true 21 | BinPackArguments: true 22 | BinPackParameters: true 23 | BraceWrapping: 24 | AfterClass: true 25 | AfterControlStatement: false 26 | AfterEnum: true 27 | AfterFunction: true 28 | AfterNamespace: false 29 | AfterObjCDeclaration: false 30 | AfterStruct: true 31 | AfterUnion: true 32 | BeforeCatch: false 33 | BeforeElse: false 34 | IndentBraces: false 35 | BreakBeforeBinaryOperators: None 36 | BreakBeforeBraces: Mozilla 37 | BreakBeforeTernaryOperators: true 38 | BreakConstructorInitializersBeforeComma: true 39 | ColumnLimit: 80 40 | CommentPragmas: '^ IWYU pragma:' 41 | ConstructorInitializerAllOnOneLineOrOnePerLine: false 42 | ConstructorInitializerIndentWidth: 2 43 | ContinuationIndentWidth: 2 44 | Cpp11BracedListStyle: false 45 | DerivePointerAlignment: false 46 | DisableFormat: false 47 | ExperimentalAutoDetectBinPacking: false 48 | ForEachMacros: [ foreach, Q_FOREACH, BOOST_FOREACH ] 49 | IncludeCategories: 50 | - Regex: '^"(llvm|llvm-c|clang|clang-c)/' 51 | Priority: 2 52 | - Regex: '^(<|"(gtest|isl|json)/)' 53 | Priority: 3 54 | - Regex: '.*' 55 | Priority: 1 56 | IndentCaseLabels: true 57 | IndentWidth: 2 58 | IndentWrappedFunctionNames: false 59 | KeepEmptyLinesAtTheStartOfBlocks: true 60 | MacroBlockBegin: '' 61 | MacroBlockEnd: '' 62 | MaxEmptyLinesToKeep: 1 63 | NamespaceIndentation: None 64 | ObjCBlockIndentWidth: 2 65 | ObjCSpaceAfterProperty: true 66 | ObjCSpaceBeforeProtocolList: false 67 | PenaltyBreakBeforeFirstCallParameter: 19 68 | PenaltyBreakComment: 300 69 | PenaltyBreakFirstLessLess: 120 70 | PenaltyBreakString: 1000 71 | PenaltyExcessCharacter: 1000000 72 | PenaltyReturnTypeOnItsOwnLine: 200 73 | PointerAlignment: Left 74 | ReflowComments: true 75 | SortIncludes: true 76 | SpaceAfterCStyleCast: false 77 | SpaceBeforeAssignmentOperators: true 78 | SpaceBeforeParens: ControlStatements 79 | SpaceInEmptyParentheses: false 80 | SpacesBeforeTrailingComments: 1 81 | SpacesInAngles: false 82 | SpacesInContainerLiterals: true 83 | SpacesInCStyleCastParentheses: false 84 | SpacesInParentheses: false 85 | SpacesInSquareBrackets: false 86 | Standard: Cpp11 87 | TabWidth: 8 88 | UseTab: Never 89 | ... 90 | 91 | -------------------------------------------------------------------------------- /neuralnet_pytorch/optim/nadam.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | from torch import optim 3 | 4 | 5 | class NAdam(optim.Adam): 6 | """ 7 | Adaptive moment with Nesterov gradients. 8 | 9 | http://cs229.stanford.edu/proj2015/054_report.pdf 10 | 11 | Parameters 12 | ---------- 13 | params 14 | iterable of parameters to optimize or dicts defining 15 | parameter groups 16 | lr 17 | learning rate (default: 1e-3) 18 | betas 19 | coefficients used for computing 20 | running averages of gradient and its square (default: (0.9, 0.999)) 21 | eps 22 | term added to the denominator to improve 23 | numerical stability (default: 1e-8) 24 | weight_decay 25 | weight decay (L2 penalty) (default: 0) 26 | decay 27 | a decay scheme for `betas[0]`. 28 | Default: :math:`\\beta * (1 - 0.5 * 0.96^{\\frac{t}{250}})` 29 | where `t` is the training step. 30 | """ 31 | 32 | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, 33 | decay=lambda x, t: x * (1. - .5 * .96 ** (t / 250.))): 34 | super().__init__(params, lr, betas, eps, weight_decay) 35 | self.decay = decay 36 | 37 | def step(self, closure=None): 38 | loss = None 39 | if closure is not None: 40 | loss = closure() 41 | 42 | for group in self.param_groups: 43 | for p in group['params']: 44 | if p.grad is None: 45 | continue 46 | grad = p.grad.data 47 | if grad.is_sparse: 48 | raise RuntimeError('NAdam does not support sparse gradients, please consider SparseAdam instead') 49 | 50 | state = self.state[p] 51 | 52 | # State initialization 53 | if len(state) == 0: 54 | state['step'] = 0 55 | # Exponential moving average of gradient values 56 | state['exp_avg'] = T.zeros_like(p.data) 57 | # Exponential moving average of squared gradient values 58 | state['exp_avg_sq'] = T.zeros_like(p.data) 59 | # Beta1 accumulation 60 | state['beta1_cum'] = 1. 61 | 62 | exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] 63 | beta1, beta2 = group['betas'] 64 | 65 | state['step'] += 1 66 | 67 | if group['weight_decay'] != 0: 68 | grad.add_(group['weight_decay'], p.data) 69 | 70 | beta1_t = self.decay(beta1, state['step']) 71 | beta1_tp1 = self.decay(beta1, state['step'] + 1.) 72 | beta1_cum = state['beta1_cum'] * beta1_t 73 | 74 | g_hat_t = grad / (1. - beta1_cum) 75 | exp_avg.mul_(beta1).add_(1. - beta1, grad) 76 | m_hat_t = exp_avg / (1. - beta1_cum * beta1_tp1) 77 | 78 | exp_avg_sq.mul_(beta2).addcmul_(1. - beta2, grad, grad) 79 | v_hat_t = exp_avg_sq / (1. - beta2 ** state['step']) 80 | m_bar_t = (1. - beta1) * g_hat_t + beta1_tp1 * m_hat_t 81 | 82 | denom = v_hat_t.sqrt().add_(group['eps']) 83 | p.data.addcdiv_(-group['lr'], m_bar_t, denom) 84 | state['beta1_cum'] = beta1_cum 85 | 86 | return loss 87 | -------------------------------------------------------------------------------- /neuralnet_pytorch/utils/numpy_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | __all__ = ['is_outlier', 'smooth'] 4 | 5 | 6 | def smooth(x, beta=.9, window='hanning'): 7 | """ 8 | Smoothens the data using a window with requested size. 9 | This method is based on the convolution of a scaled window with the signal. 10 | The signal is prepared by introducing reflected copies of the signal 11 | (with the window size) in both ends so that transient parts are minimized 12 | in the begining and end part of the output signal. 13 | 14 | :param x: 15 | the input signal. 16 | :param beta: 17 | the weighted moving average coeff. Window length is :math:`1 / (1 - \\beta)`. 18 | :param window: 19 | the type of window from ``'flat'``, ``'hanning'``, ``'hamming'``, 20 | ``'bartlett'``, and ``'blackman'``. 21 | Flat window will produce a moving average smoothing. 22 | :return: 23 | the smoothed signal. 24 | 25 | Examples 26 | -------- 27 | 28 | .. code-block:: python 29 | 30 | t = linspace(-2, 2, .1) 31 | x = sin(t) + randn(len(t)) * .1 32 | y = smooth(x) 33 | 34 | """ 35 | 36 | x = np.array(x) 37 | assert x.ndim == 1, 'smooth only accepts 1 dimension arrays' 38 | assert 0 < beta < 1, 'Input vector needs to be bigger than window size' 39 | 40 | window_len = int(1 / (1 - beta)) 41 | if window_len < 3 or x.shape[0] < window_len: 42 | return x 43 | 44 | s = np.r_[x[window_len - 1:0:-1], x, x[-2:-window_len - 1:-1]] 45 | if isinstance(window, str): 46 | assert window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman'], \ 47 | 'Window is on of \'flat\', \'hanning\', \'hamming\', \'bartlett\', \'blackman\'' 48 | 49 | if window == 'flat': # moving average 50 | w = np.ones(window_len, 'd') 51 | else: 52 | w = eval('np.' + window + '(window_len)') 53 | else: 54 | window = np.array(window) 55 | assert window.ndim == 1, 'Window must be a 1-dim array' 56 | w = window 57 | 58 | y = np.convolve(w / w.sum(), s, mode='valid') 59 | return y if y.shape[0] == x.shape[0] else y[(window_len // 2 - 1):-(window_len // 2)] 60 | 61 | 62 | def is_outlier(x: np.ndarray, thresh=3.5): 63 | """ 64 | Returns a boolean array with True if points are outliers and False 65 | otherwise. Adapted from https://stackoverflow.com/a/11886564/4591601. 66 | 67 | :param x: 68 | an ``nxd`` array of observations 69 | :param thresh: 70 | the modified z-score to use as a threshold. Observations with 71 | a modified z-score (based on the median absolute deviation) greater 72 | than this value will be classified as outliers. 73 | 74 | :return: 75 | a ``nx1`` boolean array. 76 | 77 | References 78 | ---------- 79 | Boris Iglewicz and David Hoaglin (1993), "Volume 16: How to Detect and 80 | Handle Outliers", The ASQC Basic References in Quality Control: 81 | Statistical Techniques, Edward F. Mykytka, Ph.D., Editor. 82 | """ 83 | 84 | if len(x) == 1: 85 | return np.array([False]) 86 | 87 | if len(x.shape) == 1: 88 | x = x[:, None] 89 | 90 | median = np.median(x, axis=0) 91 | diff = np.sum((x - median) ** 2, axis=-1) 92 | diff = np.sqrt(diff) 93 | med_abs_deviation = np.median(diff) 94 | modified_z_score = 0.6745 * diff / med_abs_deviation 95 | return modified_z_score > thresh 96 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from setuptools import setup, find_packages 3 | import os 4 | import versioneer 5 | 6 | cwd = os.getcwd() 7 | version_data = versioneer.get_versions() 8 | CMD_CLASS = versioneer.get_cmdclass() 9 | 10 | if version_data['error'] is not None: 11 | # Get the fallback version 12 | # We can't import neuralnet_pytorch.version as it isn't yet installed, so parse it. 13 | fname = os.path.join(os.path.split(__file__)[0], "neuralnet_pytorch", "version.py") 14 | with open(fname, "r") as f: 15 | lines = f.readlines() 16 | lines = [l for l in lines if l.startswith("FALLBACK_VERSION")] 17 | assert len(lines) == 1 18 | 19 | FALLBACK_VERSION = lines[0].split("=")[1].strip().strip('""') 20 | 21 | version_data['version'] = FALLBACK_VERSION 22 | 23 | 24 | def get_extensions(): 25 | if '--cuda-ext' in sys.argv: 26 | import glob 27 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 28 | CMD_CLASS.update({'build_ext': BuildExtension}) 29 | ext_root = os.path.join(cwd, 'neuralnet_pytorch/extensions') 30 | ext_src = glob.glob(os.path.join(ext_root, 'csrc/*.cpp')) + glob.glob(os.path.join(ext_root, 'csrc/*.cu')) 31 | ext_include = os.path.join(ext_root, 'include') 32 | sys.argv.remove("--cuda-ext") 33 | return [ 34 | CUDAExtension( 35 | name='neuralnet_pytorch.ext', 36 | sources=ext_src, 37 | include_dirs=[ext_include] 38 | )] 39 | else: 40 | return [] 41 | 42 | 43 | def setup_package(): 44 | here = os.path.abspath(os.path.dirname(__file__)) 45 | with open(os.path.join(here, 'README.md'), encoding='utf-8') as f: 46 | long_description = f.read() 47 | 48 | cuda_ext = get_extensions() 49 | setup( 50 | name='neuralnet-pytorch', 51 | version=version_data['version'], 52 | description='A high-level library on top of Pytorch.', 53 | long_description=long_description, 54 | long_description_content_type='text/markdown', 55 | url='https://github.com/justanhduc/neuralnet-pytorch', 56 | author='Duc Nguyen', 57 | author_email='adnguyen@yonsei.ac.kr', 58 | classifiers=[ 59 | 'Development Status :: 3 - Alpha', 60 | 'Intended Audience :: End Users/Desktop', 61 | 'License :: OSI Approved :: Mozilla Public License 2.0 (MPL 2.0)', 62 | 'Operating System :: Microsoft :: Windows :: Windows 10', 63 | 'Operating System :: POSIX :: Linux', 64 | 'Programming Language :: Python :: 3 :: Only', 65 | 'Programming Language :: Python :: 3.5', 66 | 'Programming Language :: Python :: 3.6', 67 | 'Programming Language :: Python :: 3.7' 68 | ], 69 | platforms=['Windows', 'Linux'], 70 | packages=find_packages(exclude=['docs', 'tests', 'examples']), 71 | ext_modules=cuda_ext, 72 | cmdclass=CMD_CLASS, 73 | install_requires=['sympy', 'matplotlib', 'scipy', 'numpy', 'imageio', 'future', 'tensorboard', 'git-python'], 74 | extras_require={ 75 | 'gin': ['gin-config'], 76 | 'geom': ['pykeops', 'geomloss'], 77 | 'visdom': ['visdom'], 78 | 'slack': ['slackclient'] 79 | }, 80 | project_urls={ 81 | 'Bug Reports': 'https://github.com/justanhduc/neuralnet-pytorch/issues', 82 | 'Source': 'https://github.com/justanhduc/neuralnet-pytorch', 83 | }, 84 | ) 85 | 86 | 87 | if __name__ == '__main__': 88 | setup_package() 89 | -------------------------------------------------------------------------------- /neuralnet_pytorch/utils/layer_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import numbers 3 | from torch._six import container_abcs 4 | import sympy as sp 5 | 6 | from . import root_logger 7 | 8 | __all__ = ['validate', 'no_dim_change_op', 'add_simple_repr', 'add_custom_repr', 'deprecated', 'get_non_none'] 9 | 10 | 11 | def _make_input_shape(m, n): 12 | def parse(x): 13 | if isinstance(x, container_abcs.Iterable): 14 | return x 15 | 16 | return sp.symbols('b:{}'.format(m), iteger=True) + (x,) + sp.symbols('x:{}'.format(n), integer=True) 17 | return parse 18 | 19 | 20 | def validate(func): 21 | """ 22 | A decorator to make sure output shape is a tuple of ``int`` s. 23 | """ 24 | 25 | def wrapper(self): 26 | shape = func(self) 27 | 28 | if shape is None: 29 | return None 30 | 31 | if isinstance(shape, numbers.Number): 32 | return int(shape) 33 | 34 | out = [None if x is None or (isinstance(x, numbers.Number) and np.isnan(x)) 35 | else int(x) if isinstance(x, numbers.Number) else x for x in shape] 36 | return tuple(out) 37 | 38 | return wrapper 39 | 40 | 41 | def no_dim_change_op(cls): 42 | """ 43 | A decorator to overwrite :attr:`~neuralnet_pytorch.layers._LayerMethod.output_shape` 44 | to an op that does not change the tensor shape. 45 | 46 | :param cls: 47 | a subclass of :class:`~neuralnet_pytorch.layers.Module`. 48 | """ 49 | 50 | @validate 51 | def output_shape(self): 52 | return None if self.input_shape is None else tuple(self.input_shape) 53 | 54 | cls.output_shape = property(output_shape) 55 | return cls 56 | 57 | 58 | def add_simple_repr(cls): 59 | """ 60 | A decorator to add a simple repr to the designated class. 61 | 62 | :param cls: 63 | a subclass of :class:`~neuralnet_pytorch.layers.Module`. 64 | """ 65 | 66 | def _repr(self): 67 | return super(cls, self).__repr__() + ' -> {}'.format(self.output_shape) 68 | 69 | setattr(cls, '__repr__', _repr) 70 | return cls 71 | 72 | 73 | def add_custom_repr(cls): 74 | """ 75 | A decorator to add a custom repr to the designated class. 76 | User should define extra_repr for the decorated class. 77 | 78 | :param cls: 79 | a subclass of :class:`~neuralnet_pytorch.layers.Module`. 80 | """ 81 | 82 | def _repr(self): 83 | return self.__class__.__name__ + '({}) -> {}'.format(self.extra_repr(), self.output_shape) 84 | 85 | setattr(cls, '__repr__', _repr) 86 | return cls 87 | 88 | 89 | def deprecated(new_func, version): 90 | def _deprecated(func): 91 | """prints out a deprecation warning""" 92 | 93 | def func_wrapper(*args, **kwargs): 94 | root_logger.warning('%s is deprecated and will be removed in version %s. Use %s instead.' % 95 | (func.__name__, version, new_func.__name__), exc_info=True) 96 | return func(*args, **kwargs) 97 | 98 | return func_wrapper 99 | return _deprecated 100 | 101 | 102 | def get_non_none(array): 103 | """ 104 | Gets the first item that is not ``None`` from the given array. 105 | 106 | :param array: 107 | an arbitrary array that is iterable. 108 | :return: 109 | the first item that is not ``None``. 110 | """ 111 | assert isinstance(array, container_abcs.Iterable) 112 | 113 | try: 114 | e = next(item for item in array if item is not None) 115 | except StopIteration: 116 | e = None 117 | return e 118 | -------------------------------------------------------------------------------- /docs/source/manual/gin.rst: -------------------------------------------------------------------------------- 1 | .. _gin: 2 | .. currentmodule:: neuralnet_pytorch 3 | 4 | 5 | ======================== 6 | :mod:`gin` -- Gin-config 7 | ======================== 8 | 9 | This page presents a list of preserved keywords for 10 | `Gin-config `_. 11 | 12 | .. contents:: Contents 13 | :depth: 2 14 | 15 | Optimizer Keywords 16 | ================== 17 | :: 18 | 19 | 'adabound' - neuralnet_pytorch.optim.AdaBound 20 | 'adadelta' - torch.optim.Adadelta 21 | 'adagrad' - torch.optim.Adagrad 22 | 'adam' - torch.optim.Adam 23 | 'adamax' - torch.optim.Adamax 24 | 'adamw' - torch.optim.AdamW 25 | 'asgd' - torch.optim.ASGD 26 | 'lbfgs' - torch.optim.LBFGS 27 | 'lookahead' - neuralnet_pytorch.optim.Lookahead 28 | 'nadam' - neuralnet_pytorch.optim.NAdam 29 | 'rmsprop' - torch.optim.RMSprop 30 | 'rprop' - torch.optim.Rprop 31 | 'sgd' - torch.optim.SGD 32 | 'sparse_adam' - torch.optim.SparseAdam 33 | 34 | If `Automatic Mixed Precision `_ is installed, 35 | there are a few extra choices coming from the package :: 36 | 37 | 'fusedadam' - apex.optimizers.FusedAdam 38 | 'fusedlamb' - apex.optimizers.FusedLAMB 39 | 'fusednovograd' - apex.optimizers.FusedNovoGrad 40 | 'fusedsgd' - apex.optimizers.FusedSGD 41 | 42 | Learning Rate Scheduler Keywords 43 | ================================ 44 | :: 45 | 46 | 'lambda_lr' - torch.optim.lr_scheduler.LambdaLR 47 | 'step_lr' - torch.optim.lr_scheduler.StepLR 48 | 'multistep_lr' - torch.optim.lr_scheduler.MultiStepLR 49 | 'exp_lr' - torch.optim.lr_scheduler.ExponentialLR 50 | 'cosine_lr' - torch.optim.lr_scheduler.CosineAnnealingLR 51 | 'plateau_lr' - torch.optim.lr_scheduler.ReduceLROnPlateau 52 | 'cyclic_lr' - torch.optim.lr_scheduler.CyclicLR 53 | 'inverse_lr' - neuralnet_pytorch.optim.lr_scheduler.InverseLR 54 | 'warm_restart' - neuralnet_pytorch.optim.lr_scheduler.WarmRestart 55 | 56 | Loss Keywords 57 | ============= 58 | :: 59 | 60 | 'l1loss' - torch.nn.L1Loss 61 | 'mseloss' - torch.nn.MSELoss 62 | 'celoss' - torch.nn.CrossEntropyLoss 63 | 'bceloss' - torch.nn.BCELoss 64 | 'bcelogit_loss' - torch.nn.BCEWithLogitsLoss 65 | 'nllloss' - torch.nn.NLLLoss 66 | 'kldloss' - torch.nn.KLDivLoss 67 | 'huberloss' - torch.nn.SmoothL1Loss 68 | 'cosineembed_loss' - torch.nn.CosineEmbeddingLoss 69 | 70 | Activation Keywords 71 | =================== 72 | :: 73 | 74 | 'elu' - torch.nn.ELU 75 | 'hardshrink' - torch.nn.Hardshrink 76 | 'hardtanh' - torch.nn.Hardtanh 77 | 'lrelu' - torch.nn.LeakyReLU 78 | 'logsig' - torch.nn.LogSigmoid 79 | 'multihead_att' - torch.nn.MultiheadAttention 80 | 'prelu' - torch.nn.PReLU 81 | 'relu' - torch.nn.ReLU 82 | 'rrelu' - torch.nn.RReLU 83 | 'selu' - torch.nn.SELU 84 | 'celu' - torch.nn.CELU 85 | 'sigmoid' - torch.nn.Sigmoid 86 | 'softplus' - torch.nn.Softplus 87 | 'softshrink' - torch.nn.Softshrink 88 | 'softsign' - torch.nn.Softsign 89 | 'tanh' - torch.nn.Tanh 90 | 'tanhshrink' - torch.nn.Tanhshrink 91 | 'threshold' - torch.nn.Threshold 92 | 93 | Data Type Keywords 94 | ================== 95 | :: 96 | 97 | 'float16' - torch.float16 98 | 'float32' - torch.float32 99 | 'float64' - torch.float64 100 | 'int8' - torch.int8 101 | 'int16' - torch.int16 102 | 'int32' - torch.int32 103 | 'int64' - torch.int64 104 | 'complex32' - torch.complex32 105 | 'complex64' - torch.complex64 106 | 'complex128' - torch.complex128 107 | 'float' - torch.float 108 | 'short' - torch.short 109 | 'long' - torch.long 110 | 'half' - torch.half 111 | 'uint8' - torch.uint8 112 | 'int' - torch.int 113 | -------------------------------------------------------------------------------- /neuralnet_pytorch/layers/adain.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | 3 | from .. import utils 4 | from .abstract import MultiMultiInputModule, MultiSingleInputModule, SingleMultiInputModule 5 | 6 | __all__ = ['AdaIN', 'MultiModuleAdaIN', 'MultiInputAdaIN'] 7 | 8 | 9 | class _AdaIN: 10 | 11 | def normalize(self, input1, input2): 12 | mean1, std1 = T.mean(input1, self.dim1, keepdim=True), T.sqrt(T.var(input1, self.dim1, keepdim=True) + 1e-8) 13 | mean2, std2 = T.mean(input2, self.dim2, keepdim=True), T.sqrt(T.var(input2, self.dim2, keepdim=True) + 1e-8) 14 | return std2 * (input1 - mean1) / std1 + mean2 15 | 16 | 17 | class AdaIN(SingleMultiInputModule, _AdaIN): 18 | """ 19 | The original Adaptive Instance Normalization from https://arxiv.org/abs/1703.06868. 20 | 21 | :math:`Y_1 = \\text{module}(X)` 22 | 23 | :math:`Y_2 = \\text{module}(X)` 24 | 25 | :math:`Y = \\sigma_{Y_2} * (Y_1 - \\mu_{Y_1}) / \\sigma_{Y_1} + \\mu_{Y_2}` 26 | 27 | Parameters 28 | ---------- 29 | module 30 | a :mod:`torch` module which generates target feature maps. 31 | dim 32 | dimension to reduce in the target feature maps. 33 | Default: ``(2, 3)``. 34 | """ 35 | 36 | def __init__(self, module, dim=(2, 3)): 37 | super().__init__(module) 38 | self.dim1 = dim 39 | self.dim2 = dim 40 | 41 | def forward(self, *input, **kwargs): 42 | out1, out2 = super().forward(*input, **kwargs) 43 | return super().normalize(out1, out2) 44 | 45 | def extra_repr(self): 46 | s = 'dim={dim1}'.format(**self.__dict__) 47 | return s 48 | 49 | 50 | class MultiModuleAdaIN(MultiSingleInputModule, _AdaIN): 51 | """ 52 | A modified Adaptive Instance Normalization from https://arxiv.org/abs/1703.06868. 53 | 54 | :math:`Y_1 = \\text{module1}(X)` 55 | 56 | :math:`Y_2 = \\text{module2}(X)` 57 | 58 | :math:`Y = \\sigma_{Y_2} * (Y_1 - \\mu_{Y_1}) / \\sigma_{Y_1} + \\mu_{Y_2}` 59 | 60 | Parameters 61 | ---------- 62 | module1 63 | a :mod:`torch` module which generates target feature maps. 64 | module2 65 | a :mod:`torch` module which generates style feature maps. 66 | dim1 67 | dimension to reduce in the target feature maps. 68 | Default: ``(2, 3)``. 69 | dim2 70 | dimension to reduce in the style feature maps. 71 | Default: ``(2, 3)``. 72 | """ 73 | 74 | def __init__(self, module1, module2, dim1=(2, 3), dim2=(2, 3)): 75 | super().__init__(module1, module2) 76 | self.dim1 = dim1 77 | self.dim2 = dim2 78 | 79 | def forward(self, input, *args, **kwargs): 80 | out1, out2 = super().forward(input, *args, **kwargs) 81 | return super().normalize(out1, out2) 82 | 83 | @property 84 | @utils.validate 85 | def output_shape(self): 86 | return self.input_shape[0] 87 | 88 | def extra_repr(self): 89 | s = 'dim1={dim1}, dim2={dim2}'.format(**self.__dict__) 90 | return s 91 | 92 | 93 | class MultiInputAdaIN(MultiMultiInputModule, _AdaIN): 94 | """ 95 | A modified Adaptive Instance Normalization from https://arxiv.org/abs/1703.06868. 96 | 97 | :math:`Y_1 = \\text{module1}(X_1)` 98 | 99 | :math:`Y_2 = \\text{module2}(X_2)` 100 | 101 | :math:`Y = \\sigma_{Y_2} * (Y_1 - \\mu_{Y_1}) / \\sigma_{Y_1} + \\mu_{Y_2}` 102 | 103 | Parameters 104 | ---------- 105 | module1 106 | a :mod:`torch` module which generates target feature maps. 107 | module2 108 | a :mod:`torch` module which generates style feature maps. 109 | dim1 110 | dimension to reduce in the target feature maps. 111 | Default: ``(2, 3)``. 112 | dim2 113 | dimension to reduce in the style feature maps. 114 | Default: ``(2, 3)``. 115 | """ 116 | 117 | def __init__(self, module1, module2, dim1=(2, 3), dim2=(2, 3)): 118 | super().__init__(module1, module2) 119 | self.dim1 = dim1 120 | self.dim2 = dim2 121 | 122 | def forward(self, *input, **kwargs): 123 | out1, out2 = super().forward(*input, **kwargs) 124 | return super().normalize(out1, out2) 125 | 126 | @property 127 | @utils.validate 128 | def output_shape(self): 129 | return self.input_shape[0] 130 | 131 | def extra_repr(self): 132 | s = 'dim1={dim1}, dim2={dim2}'.format(**self.__dict__) 133 | return s 134 | -------------------------------------------------------------------------------- /neuralnet_pytorch/optim/lookahead.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict 2 | import torch 3 | from torch.optim.optimizer import Optimizer 4 | 5 | 6 | class Lookahead(Optimizer): 7 | """ 8 | PyTorch implementation of the lookahead wrapper. 9 | Lookahead Optimizer: https://arxiv.org/abs/1907.08610. 10 | 11 | Parameters 12 | ---------- 13 | optimizer 14 | an usual optimizer such as Adam. 15 | la_steps 16 | number of lookahead steps. 17 | Default: 5. 18 | alpha 19 | linear interpolation coefficient. 20 | Default: 0.8. 21 | pullback_momentum 22 | either ``'reset'``, ``'pullback'``, or ``None``. 23 | Default: ``None``. 24 | """ 25 | 26 | def __init__(self, optimizer, la_steps=5, alpha=0.8, pullback_momentum=None): 27 | self.optimizer = optimizer 28 | self._la_step = 0 # counter for inner optimizer 29 | self.alpha = alpha 30 | self._total_la_steps = la_steps 31 | assert pullback_momentum in ['reset', 'pullback', None] 32 | self.pullback_momentum = pullback_momentum 33 | 34 | self.state = defaultdict(dict) 35 | 36 | # Cache the current optimizer parameters 37 | for group in optimizer.param_groups: 38 | for p in group['params']: 39 | param_state = self.state[p] 40 | param_state['cached_params'] = torch.zeros_like(p.data) 41 | param_state['cached_params'].copy_(p.data) 42 | if self.pullback_momentum == 'pullback': 43 | param_state['cached_mom'] = torch.zeros_like(p.data) 44 | 45 | def __getstate__(self): 46 | return { 47 | 'state': self.state, 48 | 'optimizer': self.optimizer, 49 | 'la_alpha': self.alpha, 50 | '_la_step': self._la_step, 51 | '_total_la_steps': self._total_la_steps, 52 | 'pullback_momentum': self.pullback_momentum 53 | } 54 | 55 | def zero_grad(self): 56 | self.optimizer.zero_grad() 57 | 58 | @property 59 | def la_step(self): 60 | return self._la_step 61 | 62 | @la_step.setter 63 | def la_step(self, value): 64 | self._la_step = value 65 | 66 | def state_dict(self): 67 | return self.optimizer.state_dict() 68 | 69 | def load_state_dict(self, state_dict): 70 | self.optimizer.load_state_dict(state_dict) 71 | 72 | def _backup_and_load_cache(self): 73 | for group in self.optimizer.param_groups: 74 | for p in group['params']: 75 | param_state = self.state[p] 76 | param_state['backup_params'] = torch.zeros_like(p.data) 77 | param_state['backup_params'].copy_(p.data) 78 | p.data.copy_(param_state['cached_params']) 79 | 80 | def _clear_and_load_backup(self): 81 | for group in self.optimizer.param_groups: 82 | for p in group['params']: 83 | param_state = self.state[p] 84 | p.data.copy_(param_state['backup_params']) 85 | del param_state['backup_params'] 86 | 87 | @property 88 | def param_groups(self): 89 | return self.optimizer.param_groups 90 | 91 | def step(self, closure=None): 92 | loss = self.optimizer.step(closure) 93 | self.la_step += 1 94 | 95 | if self.la_step >= self._total_la_steps: 96 | self.la_step = 0 97 | # Lookahead and cache the current optimizer parameters 98 | for group in self.optimizer.param_groups: 99 | for p in group['params']: 100 | param_state = self.state[p] 101 | p.data.mul_(self.alpha).add_(1.0 - self.alpha, param_state['cached_params']) # crucial line 102 | param_state['cached_params'].copy_(p.data) 103 | if self.pullback_momentum == 'pullback': 104 | internal_momentum = self.optimizer.state[p]['momentum_buffer'] 105 | self.optimizer.state[p]['momentum_buffer'] = internal_momentum.mul_(self.alpha).add_( 106 | 1.0 - self.alpha, param_state['cached_mom']) 107 | param_state['cached_mom'] = self.optimizer.state[p]['momentum_buffer'] 108 | elif self.pullback_momentum == 'reset': 109 | self.optimizer.state[p]['momentum_buffer'] = torch.zeros_like(p.data) 110 | 111 | return loss 112 | -------------------------------------------------------------------------------- /neuralnet_pytorch/layers/aggregation.py: -------------------------------------------------------------------------------- 1 | from .abstract import Module, MultiSingleInputModule, MultiMultiInputModule 2 | from .. import utils 3 | 4 | __all__ = ['Activation', 'Sum', 'SequentialSum', 'ConcurrentSum'] 5 | 6 | 7 | @utils.add_simple_repr 8 | @utils.no_dim_change_op 9 | class Activation(Module): 10 | """ 11 | Applies a non-linear function to the incoming input. 12 | 13 | Parameters 14 | ---------- 15 | activation 16 | non-linear function to activate the linear result. 17 | It accepts any callable function 18 | as well as a recognizable ``str``. 19 | A list of possible ``str`` is in :func:`~neuralnet_pytorch.utils.function`. 20 | input_shape 21 | shape of the input tensor. Can be ``None``. 22 | kwargs 23 | extra keyword arguments to pass to activation. 24 | """ 25 | 26 | def __init__(self, activation='relu', input_shape=None, **kwargs): 27 | super().__init__(input_shape) 28 | self.activation = utils.function(activation, **kwargs) 29 | 30 | def forward(self, input, *args, **kwargs): 31 | return self.activation(input) 32 | 33 | def extra_repr(self): 34 | s = 'activation={}'.format(self.activation.__name__) 35 | return s 36 | 37 | 38 | class Sum(MultiSingleInputModule): 39 | """ 40 | Sums the outputs of multiple modules given an input tensor. 41 | A subclass of :class:`~neuralnet_pytorch.layers.MultiSingleInputModule`. 42 | 43 | See Also 44 | -------- 45 | :class:`~neuralnet_pytorch.layers.MultiSingleInputModule` 46 | :class:`~neuralnet_pytorch.layers.MultiMultiInputModule` 47 | :class:`~neuralnet_pytorch.layers.SequentialSum` 48 | :class:`~neuralnet_pytorch.layers.ConcurrentSum` 49 | :class:`~neuralnet_pytorch.resizing.Cat` 50 | :class:`~neuralnet_pytorch.resizing.SequentialCat` 51 | :class:`~neuralnet_pytorch.resizing.ConcurrentCat` 52 | """ 53 | 54 | def __init__(self, *modules_or_tensors): 55 | super().__init__(*modules_or_tensors) 56 | 57 | def forward(self, input, *args, **kwargs): 58 | outputs = super().forward(input) 59 | return sum(outputs) 60 | 61 | @property 62 | @utils.validate 63 | def output_shape(self): 64 | if None in self.input_shape: 65 | return None 66 | 67 | shapes_transposed = [item for item in zip(*self.input_shape)] 68 | input_shape_none_filtered = list(map(utils.get_non_none, shapes_transposed)) 69 | 70 | return tuple(input_shape_none_filtered) 71 | 72 | 73 | class SequentialSum(Sum): 74 | """ 75 | Sums the intermediate outputs of multiple sequential modules given an input tensor. 76 | A subclass of :class:`~neuralnet_pytorch.layers.Sum`. 77 | 78 | See Also 79 | -------- 80 | :class:`~neuralnet_pytorch.layers.MultiSingleInputModule` 81 | :class:`~neuralnet_pytorch.layers.MultiMultiInputModule` 82 | :class:`~neuralnet_pytorch.layers.Sum` 83 | :class:`~neuralnet_pytorch.layers.ConcurrentSum` 84 | :class:`~neuralnet_pytorch.resizing.Cat` 85 | :class:`~neuralnet_pytorch.resizing.SequentialCat` 86 | :class:`~neuralnet_pytorch.resizing.ConcurrentCat` 87 | """ 88 | 89 | def __init__(self, *modules): 90 | super().__init__(*modules) 91 | 92 | def forward(self, input, *args, **kwargs): 93 | outputs = [] 94 | output = input 95 | for name, module in self.named_children(): 96 | if name.startswith('tensor'): 97 | outputs.append(module()) 98 | else: 99 | output = module(output) 100 | outputs.append(output) 101 | 102 | return sum(outputs) 103 | 104 | 105 | class ConcurrentSum(MultiMultiInputModule): 106 | """ 107 | Sums the outputs of multiple modules given input tensors. 108 | A subclass of :class:`~neuralnet_pytorch.layers.MultiMultiInputModule`. 109 | 110 | See Also 111 | -------- 112 | :class:`~neuralnet_pytorch.layers.MultiSingleInputModule` 113 | :class:`~neuralnet_pytorch.layers.MultiMultiInputModule` 114 | :class:`~neuralnet_pytorch.layers.Sum` 115 | :class:`~neuralnet_pytorch.layers.SequentialSum` 116 | :class:`~neuralnet_pytorch.resizing.Cat` 117 | :class:`~neuralnet_pytorch.resizing.SequentialCat` 118 | :class:`~neuralnet_pytorch.resizing.ConcurrentCat` 119 | """ 120 | 121 | def __init__(self, *modules_or_tensors): 122 | super().__init__(*modules_or_tensors) 123 | 124 | def forward(self, *input, **kwargs): 125 | outputs = super().forward(*input) 126 | return sum(outputs) 127 | 128 | @property 129 | @utils.validate 130 | def output_shape(self): 131 | if None in self.input_shape: 132 | return None 133 | 134 | shapes_transposed = [item for item in zip(*self.input_shape)] 135 | input_shape_none_filtered = list(map(utils.get_non_none, shapes_transposed)) 136 | 137 | return tuple(input_shape_none_filtered) 138 | -------------------------------------------------------------------------------- /neuralnet_pytorch/optim/adabound.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | import math 3 | from torch import optim 4 | 5 | 6 | class AdaBound(optim.Optimizer): 7 | """ 8 | Implements AdaBound algorithm proposed in 9 | `Adaptive Gradient Methods with Dynamic Bound of Learning Rate`_. 10 | 11 | .. _Adaptive Gradient Methods with Dynamic Bound of Learning Rate: 12 | https://openreview.net/forum?id=Bkg3g2R9FX 13 | 14 | Parameters 15 | ---------- 16 | params 17 | iterable of parameters to optimize or dicts defining. 18 | parameter groups 19 | lr 20 | Adam learning rate. Default: 1e-3. 21 | betas 22 | coefficients used for computing running averages of gradient 23 | and its square. Default: (0.9, 0.999). 24 | final_lr 25 | final (SGD) learning rate. Default: 0.1. 26 | gamma 27 | convergence speed of the bound functions. Default: 1e-3. 28 | eps 29 | term added to the denominator to improve 30 | numerical stability. Default: 1e-8. 31 | weight_decay 32 | weight decay (L2 penalty). Default: 0. 33 | amsbound : bool 34 | whether to use the AMSBound variant of this algorithm. 35 | """ 36 | 37 | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), final_lr=0.1, gamma=1e-3, 38 | eps=1e-8, weight_decay=0, amsbound=False): 39 | if not 0.0 <= lr: 40 | raise ValueError("Invalid learning rate: {}".format(lr)) 41 | if not 0.0 <= eps: 42 | raise ValueError("Invalid epsilon value: {}".format(eps)) 43 | if not 0.0 <= betas[0] < 1.0: 44 | raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) 45 | if not 0.0 <= betas[1] < 1.0: 46 | raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) 47 | if not 0.0 <= final_lr: 48 | raise ValueError("Invalid final learning rate: {}".format(final_lr)) 49 | if not 0.0 <= gamma < 1.0: 50 | raise ValueError("Invalid gamma parameter: {}".format(gamma)) 51 | defaults = dict(lr=lr, betas=betas, final_lr=final_lr, gamma=gamma, eps=eps, 52 | weight_decay=weight_decay, amsbound=amsbound) 53 | super(AdaBound, self).__init__(params, defaults) 54 | 55 | self.base_lrs = list(map(lambda group: group['lr'], self.param_groups)) 56 | 57 | def __setstate__(self, state): 58 | super(AdaBound, self).__setstate__(state) 59 | for group in self.param_groups: 60 | group.setdefault('amsbound', False) 61 | 62 | def step(self, closure=None): 63 | loss = None 64 | if closure is not None: 65 | loss = closure() 66 | 67 | for group, base_lr in zip(self.param_groups, self.base_lrs): 68 | for p in group['params']: 69 | if p.grad is None: 70 | continue 71 | grad = p.grad.data 72 | if grad.is_sparse: 73 | raise RuntimeError( 74 | 'Adam does not support sparse gradients, please consider SparseAdam instead') 75 | amsbound = group['amsbound'] 76 | 77 | state = self.state[p] 78 | 79 | # State initialization 80 | if len(state) == 0: 81 | state['step'] = 0 82 | # Exponential moving average of gradient values 83 | state['exp_avg'] = T.zeros_like(p.data) 84 | # Exponential moving average of squared gradient values 85 | state['exp_avg_sq'] = T.zeros_like(p.data) 86 | if amsbound: 87 | # Maintains max of all exp. moving avg. of sq. grad. values 88 | state['max_exp_avg_sq'] = T.zeros_like(p.data) 89 | 90 | exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] 91 | if amsbound: 92 | max_exp_avg_sq = state['max_exp_avg_sq'] 93 | beta1, beta2 = group['betas'] 94 | 95 | state['step'] += 1 96 | 97 | if group['weight_decay'] != 0: 98 | grad = grad.add(group['weight_decay'], p.data) 99 | 100 | # Decay the first and second moment running average coefficient 101 | exp_avg.mul_(beta1).add_(1 - beta1, grad) 102 | exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) 103 | if amsbound: 104 | # Maintains the maximum of all 2nd moment running avg. till now 105 | T.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq) 106 | # Use the max. for normalizing running avg. of gradient 107 | denom = max_exp_avg_sq.sqrt().add_(group['eps']) 108 | else: 109 | denom = exp_avg_sq.sqrt().add_(group['eps']) 110 | 111 | bias_correction1 = 1 - beta1 ** state['step'] 112 | bias_correction2 = 1 - beta2 ** state['step'] 113 | step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1 114 | 115 | # Applies bounds on actual learning rate 116 | # lr_scheduler cannot affect final_lr, this is a workaround to apply lr decay 117 | final_lr = group['final_lr'] * group['lr'] / base_lr 118 | lower_bound = final_lr * (1 - 1 / (group['gamma'] * state['step'] + 1)) 119 | upper_bound = final_lr * (1 + 1 / (group['gamma'] * state['step'])) 120 | step_size = T.full_like(denom, step_size) 121 | step_size.div_(denom).clamp_(lower_bound, upper_bound).mul_(exp_avg) 122 | 123 | p.data.add_(-step_size) 124 | 125 | return loss 126 | -------------------------------------------------------------------------------- /tests/misc_test.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | import numpy as np 3 | from torch import testing 4 | import pytest 5 | 6 | import neuralnet_pytorch as nnt 7 | from neuralnet_pytorch import cuda_available 8 | 9 | dev = ('cpu', 'cuda') if cuda_available else ('cpu',) 10 | 11 | 12 | def test_gin_config(): 13 | try: 14 | from neuralnet_pytorch import gin_nnt as gin 15 | except ImportError: 16 | print('Please install Gin-config first and run this test again') 17 | return 18 | 19 | @gin.configurable('net') 20 | def assert_same(dtype, activation, loss, optimizer, scheduler): 21 | assert dtype is T.float32 22 | assert isinstance(activation(), T.nn.Tanh) 23 | assert isinstance(loss(), T.nn.L1Loss) 24 | 25 | import os 26 | config_file = os.path.join(os.path.dirname(__file__), 'test_files/test.gin') 27 | gin.parse_config_file(config_file) 28 | assert_same() 29 | 30 | 31 | @pytest.mark.parametrize('device', dev) 32 | def test_track(device): 33 | shape = (2, 3, 5, 5) 34 | a = T.rand(*shape).to(device) 35 | 36 | conv1 = nnt.track('op', nnt.Conv2d(shape, 4, 3), 'all').to(device) 37 | conv2 = nnt.Conv2d(conv1.output_shape, 5, 3).to(device) 38 | intermediate = conv1(a) 39 | output = nnt.track('conv2_output', conv2(intermediate), 'all').to(device) 40 | loss = T.sum(output ** 2) 41 | loss.backward(retain_graph=True) 42 | d_inter = T.autograd.grad(loss, intermediate, retain_graph=True) 43 | d_out = T.autograd.grad(loss, output) 44 | tracked = nnt.eval_tracked_variables() 45 | 46 | testing.assert_allclose(tracked['conv2_output'], nnt.utils.to_numpy(output)) 47 | testing.assert_allclose(np.stack(tracked['grad_conv2_output']), nnt.utils.to_numpy(d_out[0])) 48 | testing.assert_allclose(tracked['op'], nnt.utils.to_numpy(intermediate)) 49 | for d_inter_, tracked_d_inter_ in zip(d_inter, tracked['grad_op_output']): 50 | testing.assert_allclose(tracked_d_inter_, nnt.utils.to_numpy(d_inter_)) 51 | 52 | 53 | @pytest.mark.parametrize('device', dev) 54 | def test_monitor(device): 55 | shape = (64, 512) 56 | a = T.rand(*shape).to(device) 57 | 58 | n_epochs = 5 59 | n_iters = 10 60 | print_freq = 5 61 | 62 | mon = nnt.Monitor(use_tensorboard=True, print_freq=print_freq) 63 | mon.dump('foo.pkl', a) 64 | loaded = mon.load('foo.pkl') 65 | testing.assert_allclose(a, loaded) 66 | 67 | mon.dump('foo.txt', nnt.utils.to_numpy(a), 'txt') 68 | loaded = T.from_numpy(mon.load('foo.txt', 'txt', dtype='float32')).to(device) 69 | testing.assert_allclose(a, loaded) 70 | 71 | mon.dump('foo.pt', {'a': a}, 'torch') 72 | loaded = mon.load('foo.pt', 'torch')['a'] 73 | testing.assert_allclose(a, loaded) 74 | 75 | mon.epoch = 0 76 | for epoch in mon.iter_epoch(range(n_epochs)): 77 | for it in mon.iter_batch(range(n_iters)): 78 | mon.plot('parabol', (it + epoch) ** 2) 79 | mon.hist('histogram', a + (it * epoch)) 80 | mon.imwrite('image', a[None, None]) 81 | mon.plot_matrix('random', a) 82 | mon.dump('foo.pkl', a + it, keep=3) 83 | mon.dump('foo.txt', nnt.utils.to_numpy(a + it), 'txt', keep=3) 84 | mon.dump('foo.pt', {'a': a + it}, 'torch', keep=4) 85 | 86 | loaded = mon.load('foo.pkl', version=48) 87 | testing.assert_allclose(a + 8., loaded) 88 | loaded = T.from_numpy(mon.load('foo.txt', 'txt', version=48, dtype='float32')).to(device) 89 | testing.assert_allclose(a + 8., loaded) 90 | loaded = mon.load('foo.pt', 'torch', version=49)['a'] 91 | testing.assert_allclose(a + 9, loaded) 92 | 93 | 94 | @pytest.mark.parametrize( 95 | 'idx', 96 | (slice(None, None), slice(None, 2), slice(1, None), slice(1, 3)) 97 | ) 98 | def test_slicing_sequential(idx): 99 | input_shape = (None, 3, 256, 256) 100 | 101 | a = nnt.Sequential(input_shape=input_shape) 102 | a.conv1 = nnt.Conv2d(a.output_shape, 64, 3) 103 | a.conv2 = nnt.Conv2d(a.output_shape, 128, 3) 104 | a.conv3 = nnt.Conv2d(a.output_shape, 256, 3) 105 | a.conv4 = nnt.Conv2d(a.output_shape, 512, 3) 106 | 107 | b = a[idx] 108 | start = 0 if idx.start is None else idx.start 109 | assert b.input_shape == a[start].input_shape 110 | 111 | class Foo(nnt.Sequential): 112 | def __init__(self, input_shape): 113 | super().__init__(input_shape=input_shape) 114 | 115 | self.conv1 = nnt.Conv2d(self.output_shape, 64, 3) 116 | self.conv2 = nnt.Conv2d(self.output_shape, 128, 3) 117 | self.conv3 = nnt.Conv2d(self.output_shape, 256, 3) 118 | self.conv4 = nnt.Conv2d(self.output_shape, 512, 3) 119 | 120 | foo = Foo(input_shape) 121 | b = foo[idx] 122 | start = 0 if idx.start is None else idx.start 123 | assert isinstance(b, nnt.Sequential) 124 | assert b.input_shape == a[start].input_shape 125 | 126 | 127 | @pytest.mark.parametrize('device', dev) 128 | @pytest.mark.parametrize('bs', (pytest.param(None, marks=pytest.mark.xfail), 1, 2, 3, 4, 5)) 129 | @pytest.mark.parametrize('shuffle', (True, False)) 130 | @pytest.mark.parametrize('drop_last', (True, False)) 131 | @pytest.mark.parametrize('pin_memory', (True, False)) 132 | def test_dataloader(device, bs, shuffle, drop_last, pin_memory): 133 | if device == 'cpu' and pin_memory: 134 | pytest.skip('Wrong settings for pin_memory=True and device=cpu') 135 | 136 | from torch.utils.data import TensorDataset 137 | data, label = T.arange(10), T.arange(10) + 10 138 | dataset = TensorDataset(data, label) 139 | loader = nnt.DataLoader(dataset, batch_size=bs, shuffle=shuffle, drop_last=drop_last, pin_memory=pin_memory, 140 | num_workers=10) 141 | if 'cuda' in device: 142 | loader = nnt.DataPrefetcher(loader) 143 | 144 | for epoch in range(2): 145 | data_, label_ = [], [] 146 | for batch in loader: 147 | data_.append(batch[0]) 148 | label_.append(batch[1]) 149 | 150 | data_ = sorted(T.cat(data_)) 151 | label_ = sorted(T.cat(label_)) 152 | if len(data_) != len(data): 153 | assert len(data_) == len(data) // bs * bs 154 | assert len(label_) == len(label) // bs * bs 155 | else: 156 | testing.assert_allclose(data_, data) 157 | testing.assert_allclose(label_, label) 158 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Introduction 2 | ![Python - Version](https://img.shields.io/pypi/pyversions/neuralnet-pytorch.svg) 3 | [![PyPI - Version](https://img.shields.io/pypi/v/neuralnet-pytorch.svg)](https://pypi.org/project/neuralnet-pytorch/) 4 | [![PyPI - Wheel](https://img.shields.io/pypi/wheel/neuralnet-pytorch.svg)](https://pypi.org/project/neuralnet-pytorch/) 5 | [![Github - Tag](https://img.shields.io/github/tag/justanhduc/neuralnet-pytorch.svg)](https://github.com/justanhduc/neuralnet-pytorch/releases/tag/rel-0.0.4) 6 | [![License](https://img.shields.io/github/license/justanhduc/neuralnet-pytorch.svg)](https://github.com/justanhduc/neuralnet-pytorch/blob/master/LICENSE.txt) 7 | [![Build Status](https://travis-ci.org/justanhduc/neuralnet-pytorch.svg?branch=master)](https://travis-ci.org/justanhduc/neuralnet-pytorch) 8 | [![Documentation Status](https://readthedocs.org/projects/neuralnet-pytorch/badge/?version=latest)](https://neuralnet-pytorch.readthedocs.io/en/latest/?badge=latest) 9 | [![codecov](https://codecov.io/gh/justanhduc/neuralnet-pytorch/branch/master/graph/badge.svg)](https://codecov.io/gh/justanhduc/neuralnet-pytorch) 10 | 11 | __A high level framework for general purpose neural networks in Pytorch.__ 12 | 13 | Personally, going from Theano to Pytorch is pretty much like 14 | time traveling from 90s to the modern day. 15 | However, we feel like despite having a lot of bells and whistles, 16 | Pytorch is still missing many elements 17 | that are confirmed to never be added to the library. 18 | Therefore, this library is written to supplement more features 19 | to the current magical Pytorch. 20 | All the modules in the package directly subclass 21 | the corresponding modules from Pytorch, 22 | so everything should still be familiar. 23 | 24 | # At first glance 25 | Neuralnet-pytorch mostly preserves the same spirit of native Pytorch but in a 26 | (hopefully) less verbose way. 27 | The most prominent feature of Neuralnet-pytorch is the ability to handle 28 | input and output tensor shapes of operators abstractly 29 | (powered by [Sympy](https://docs.sympy.org/latest/index.html)). 30 | For example, the following snippet in Pytorch 31 | 32 | ``` 33 | >>> from torch import nn 34 | >>> model = nn.Sequential( 35 | ... nn.Conv2d(1, 20, 5, padding=0), 36 | ... nn.ReLU(), 37 | ... nn.Conv2d(20, 64, 5, padding=0), 38 | ... nn.ReLU() 39 | ... ) 40 | >>> print(model) 41 | Sequential( 42 | (0): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1)) 43 | (1): ReLU() 44 | (2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1)) 45 | (3): ReLU() 46 | ) 47 | ``` 48 | 49 | can be rewritten in Neuralnet-pytorch as 50 | ``` 51 | >>> import neuralnet_pytorch as nnt 52 | >>> model = nnt.Sequential( 53 | ... nnt.Conv2d(1, 20, 5, padding=0, activation='relu'), 54 | ... nnt.Conv2d(20, 64, 5, padding=0, activation='relu') 55 | ... ) 56 | >>> print(model) 57 | Sequential( 58 | (0): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1), activation=relu) -> (b0, 20, x0 - 4, x1 - 4) 59 | (1): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1), activation=relu) -> (b0, 64, x0 - 8, x1 - 8) 60 | ) -> (b0, 64, x0 - 8, x1 - 8) 61 | 62 | ``` 63 | 64 | which is the same as the native Pytorch, or 65 | 66 | ``` 67 | >>> import neuralnet_pytorch as nnt 68 | >>> model = nnt.Sequential(input_shape=1) 69 | >>> model.conv1 = nnt.Conv2d(model.output_shape, 20, 5, padding=0, activation='relu') 70 | >>> model.conv2 = nnt.Conv2d(model.output_shape, 64, 5, padding=0, activation='relu') 71 | >>> print(model) 72 | Sequential( 73 | (conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1), activation=relu) -> (b0, 20, x0 - 4, x1 - 4) 74 | (conv2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1), activation=relu) -> (b0, 64, x0 - 8, x1 - 8) 75 | ) -> (b0, 64, x0 - 8, x1 - 8) 76 | ``` 77 | 78 | which frees you from a lot of memorization and manual calculations when adding one layer on top of another. 79 | Theano folks will also find some reminiscence as many functions are highly inspired by Theano. 80 | 81 | # Requirements 82 | 83 | [Pytorch](https://pytorch.org/) >= 1.0.0 84 | 85 | [Sympy](https://docs.sympy.org/latest/index.html) 86 | 87 | [Matplotlib](https://matplotlib.org/) 88 | 89 | [Visdom](https://github.com/facebookresearch/visdom) 90 | 91 | [Tensorboard](https://www.tensorflow.org/tensorboard) 92 | 93 | [Gin-config](https://github.com/google/gin-config) (optional) 94 | 95 | [Pykeops](https://github.com/getkeops/keops) 96 | (optional, required for the calculation of Sinkhorn Wasserstein loss.) 97 | 98 | [Geomloss](https://github.com/jeanfeydy/geomloss) 99 | (optional, required for the calculation of Sinkhorn Wasserstein loss.) 100 | 101 | # Documentation 102 | 103 | https://neuralnet-pytorch.readthedocs.io 104 | 105 | # Installation 106 | 107 | Stable version 108 | ``` 109 | pip install --upgrade neuralnet-pytorch 110 | ``` 111 | 112 | Bleeding-edge version (recommended) 113 | 114 | ``` 115 | pip install git+git://github.com/justanhduc/neuralnet-pytorch.git@master 116 | ``` 117 | 118 | To install the package with optional dependencies, try 119 | 120 | ``` 121 | pip install "neuralnet-pytorch[option] @ git+git://github.com/justanhduc/neuralnet-pytorch.git@master" 122 | ``` 123 | in which `option` can be `gin`/`geom`/`visdom`/`slack`. 124 | 125 | To install the version with some collected Cuda/C++ ops, use 126 | 127 | ``` 128 | pip install git+git://github.com/justanhduc/neuralnet-pytorch.git@master --global-option="--cuda-ext" 129 | ``` 130 | 131 | # Usages 132 | 133 | The manual reference is still under development and is available at https://neuralnet-pytorch.readthedocs.io. 134 | 135 | # TODO 136 | 137 | - [x] Adding introduction and installation 138 | - [x] Writing documentations 139 | - [x] Adding examples 140 | - [ ] Adding benchmarks for examples 141 | 142 | # Disclaimer 143 | 144 | This package is a product from my little free time during my PhD, 145 | so most but not all the written modules are properly checked. 146 | No replacements or refunds for buggy performance. 147 | All PRs are welcome. 148 | 149 | # Acknowledgements 150 | 151 | The CUDA Chamfer distance is taken from the [AtlasNet](https://github.com/ThibaultGROUEIX/AtlasNet) repo. 152 | 153 | The AdaBound optimizer is taken from its [official repo](https://github.com/Luolc/AdaBound). 154 | 155 | The adapted Gin for Pytorch code is taken from [Gin-config](https://github.com/google/gin-config). 156 | 157 | The monitor scheme is inspired from [WGAN](https://github.com/igul222/improved_wgan_training). 158 | 159 | The EMD CUDA implementation is adapted form [here](https://github.com/daerduoCarey/PyTorchEMD). 160 | -------------------------------------------------------------------------------- /docs/source/conf.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # 3 | # Configuration file for the Sphinx documentation builder. 4 | # 5 | # This file does only contain a selection of the most common options. For a 6 | # full list see the documentation: 7 | # http://www.sphinx-doc.org/en/master/config 8 | 9 | # -- Path setup -------------------------------------------------------------- 10 | 11 | # If extensions (or modules to document with autodoc) are in another directory, 12 | # add these directories to sys.path here. If the directory is relative to the 13 | # documentation root, use os.path.abspath to make it absolute, like shown here. 14 | # 15 | import os 16 | import sys 17 | sys.path.insert(0, os.path.abspath('../..')) 18 | import versioneer 19 | from neuralnet_pytorch import zoo 20 | 21 | autodoc_mock_imports = ['torch', 'numpy', 'visdom', 'matplotlib', 'scipy', 'slackclient', 'tb-nightly', 'imageio'] 22 | 23 | # -- Project information ----------------------------------------------------- 24 | 25 | project = 'Neuralnet-pytorch' 26 | copyright = '2020, Duc Nguyen' 27 | author = 'Duc Nguyen' 28 | 29 | # From Theano 30 | # We need this hokey-pokey because versioneer needs the current 31 | # directory to be the root of the project to work. 32 | _curpath = os.getcwd() 33 | os.chdir(os.path.join(os.path.dirname(os.path.dirname(__file__)), os.pardir)) 34 | # The full version, including alpha/beta/rc tags. 35 | release = versioneer.get_version() 36 | # The short X.Y version. 37 | version = '.'.join(release.split('.')[:2]) 38 | os.chdir(_curpath) 39 | del _curpath 40 | 41 | # -- General configuration --------------------------------------------------- 42 | 43 | # If your documentation needs a minimal Sphinx version, state it here. 44 | # 45 | # needs_sphinx = '1.0' 46 | 47 | # Add any Sphinx extension module names here, as strings. They can be 48 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 49 | # ones. 50 | extensions = [ 51 | 'sphinx.ext.autosectionlabel', 52 | 'sphinx.ext.autodoc', 53 | 'sphinx.ext.napoleon', 54 | 'sphinx.ext.coverage', 55 | 'sphinx.ext.viewcode' 56 | ] 57 | 58 | napoleon_google_docstring = True 59 | napoleon_numpy_docstring = True 60 | napoleon_include_init_with_doc = False 61 | napoleon_include_private_with_doc = False 62 | napoleon_include_special_with_doc = True 63 | napoleon_use_admonition_for_examples = False 64 | napoleon_use_admonition_for_notes = False 65 | napoleon_use_admonition_for_references = False 66 | napoleon_use_ivar = False 67 | napoleon_use_param = True 68 | napoleon_use_rtype = True 69 | 70 | # Add any paths that contain templates here, relative to this directory. 71 | templates_path = ['_templates'] 72 | 73 | # The suffix(es) of source filenames. 74 | # You can specify multiple suffix as a list of string: 75 | # 76 | # source_suffix = ['.rst', '.md'] 77 | source_suffix = '.rst' 78 | 79 | # The master toctree document. 80 | master_doc = 'index' 81 | 82 | # The language for content autogenerated by Sphinx. Refer to documentation 83 | # for a list of supported languages. 84 | # 85 | # This is also used if you do content translation via gettext catalogs. 86 | # Usually you set "language" from the command line for these cases. 87 | language = None 88 | 89 | # List of patterns, relative to source directory, that match files and 90 | # directories to ignore when looking for source files. 91 | # This pattern also affects html_static_path and html_extra_path. 92 | exclude_patterns = [] 93 | 94 | # The name of the Pygments (syntax highlighting) style to use. 95 | pygments_style = None 96 | 97 | # no pdf output 98 | sphinx_enable_pdf_build = False 99 | 100 | 101 | # -- Options for HTML output ------------------------------------------------- 102 | 103 | # The theme to use for HTML and HTML Help pages. See the documentation for 104 | # a list of builtin themes. 105 | # 106 | html_theme = 'sphinx_rtd_theme' 107 | html_logo = '_static/nnt-logo.png' 108 | # Theme options are theme-specific and customize the look and feel of a theme 109 | # further. For a list of options available for each theme, see the 110 | # documentation. 111 | # 112 | # html_theme_options = {} 113 | 114 | # Add any paths that contain custom static files (such as style sheets) here, 115 | # relative to this directory. They are copied after the builtin static files, 116 | # so a file named "default.css" will overwrite the builtin "default.css". 117 | html_static_path = ['_static'] 118 | 119 | # Custom sidebar templates, must be a dictionary that maps document names 120 | # to template names. 121 | # 122 | # The default sidebars (for documents that don't match any pattern) are 123 | # defined by theme itself. Builtin themes are using these templates by 124 | # default: ``['localtoc.html', 'relations.html', 'sourcelink.html', 125 | # 'searchbox.html']``. 126 | # 127 | # html_sidebars = {} 128 | 129 | html_show_sourcelink = True 130 | 131 | # -- Options for HTMLHelp output --------------------------------------------- 132 | 133 | # Output file base name for HTML help builder. 134 | htmlhelp_basename = 'neuralnet-pytorchdoc' 135 | 136 | 137 | # -- Options for LaTeX output ------------------------------------------------ 138 | 139 | latex_elements = { 140 | # The paper size ('letterpaper' or 'a4paper'). 141 | # 142 | # 'papersize': 'letterpaper', 143 | 144 | # The font size ('10pt', '11pt' or '12pt'). 145 | # 146 | # 'pointsize': '10pt', 147 | 148 | # Additional stuff for the LaTeX preamble. 149 | # 150 | # 'preamble': '', 151 | 152 | # Latex figure (float) alignment 153 | # 154 | # 'figure_align': 'htbp', 155 | } 156 | 157 | # Grouping the document tree into LaTeX files. List of tuples 158 | # (source start file, target name, title, 159 | # author, documentclass [howto, manual, or own class]). 160 | latex_documents = [ 161 | (master_doc, 'neuralnet-pytorch.tex', 'neuralnet-pytorch Documentation', 162 | 'Duc Nguyen', 'manual'), 163 | ] 164 | 165 | 166 | # -- Options for manual page output ------------------------------------------ 167 | 168 | # One entry per manual page. List of tuples 169 | # (source start file, name, description, authors, manual section). 170 | man_pages = [ 171 | (master_doc, 'neuralnet-pytorch', 'neuralnet-pytorch Documentation', 172 | [author], 1) 173 | ] 174 | 175 | 176 | # -- Options for Texinfo output ---------------------------------------------- 177 | 178 | # Grouping the document tree into Texinfo files. List of tuples 179 | # (source start file, target name, title, author, 180 | # dir menu entry, description, category) 181 | texinfo_documents = [ 182 | (master_doc, 'neuralnet-pytorch', 'neuralnet-pytorch Documentation', 183 | author, 'neuralnet-pytorch', 'One line description of project.', 184 | 'Miscellaneous'), 185 | ] 186 | 187 | 188 | # -- Options for Epub output ------------------------------------------------- 189 | 190 | # Bibliographic Dublin Core info. 191 | epub_title = project 192 | 193 | # The unique identifier of the text. This can be a ISBN number 194 | # or the project homepage. 195 | # 196 | # epub_identifier = '' 197 | 198 | # A unique identification for the text. 199 | # 200 | # epub_uid = '' 201 | 202 | # A list of files that should not be packed into the epub file. 203 | epub_exclude_files = ['search.html'] -------------------------------------------------------------------------------- /neuralnet_pytorch/gin_nnt/external_configurables.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | import torch.optim as optim 3 | import torch.nn as nn 4 | from gin import config 5 | 6 | from .. import optim as nnt_optim 7 | from .. import zoo 8 | 9 | # optimizers 10 | config.external_configurable(optim.SGD, 'sgd', module='T.optim') 11 | config.external_configurable(optim.Adam, 'adam', module='T.optim') 12 | config.external_configurable(optim.Adadelta, 'adadelta', module='T.optim') 13 | config.external_configurable(optim.Adagrad, 'adagrad', module='T.optim') 14 | config.external_configurable(optim.SparseAdam, 'sparse_adam', module='T.optim') 15 | config.external_configurable(optim.Adamax, 'adamax', module='T.optim') 16 | config.external_configurable(optim.AdamW, 'adamw', module='T.optim') 17 | config.external_configurable(optim.ASGD, 'asgd', module='T.optim') 18 | config.external_configurable(optim.LBFGS, 'lbfgs', module='T.optim') 19 | config.external_configurable(optim.RMSprop, 'rmsprop', module='T.optim') 20 | config.external_configurable(optim.Rprop, 'rprop', module='T.optim') 21 | config.external_configurable(nnt_optim.AdaBound, 'adabound', module='nnt') 22 | config.external_configurable(nnt_optim.Lookahead, 'lookahead', module='nnt') 23 | config.external_configurable(nnt_optim.NAdam, 'nadam', module='nnt') 24 | 25 | try: 26 | import apex 27 | config.external_configurable(apex.optimizers.FusedAdam, 'fusedadam', module='apex.optimizers') 28 | config.external_configurable(apex.optimizers.FusedSGD, 'fusedsgd', module='apex.optimizers') 29 | config.external_configurable(apex.optimizers.FusedNovoGrad, 'fusednovograd', module='apex.optimizers') 30 | config.external_configurable(apex.optimizers.FusedLAMB, 'fusedlamb', module='apex.optimizers') 31 | print('Apex Fused Optimizers is availble for GIN') 32 | except ImportError: 33 | pass 34 | 35 | # lr scheduler 36 | config.external_configurable(optim.lr_scheduler.LambdaLR, 'lambda_lr', module='T.optim.lr_scheduler') 37 | config.external_configurable(optim.lr_scheduler.StepLR, 'step_lr', module='T.optim.lr_scheduler') 38 | config.external_configurable(optim.lr_scheduler.MultiStepLR, 'multistep_lr', module='T.optim.lr_scheduler') 39 | config.external_configurable(optim.lr_scheduler.ExponentialLR, 'exp_lr', module='T.optim.lr_scheduler') 40 | config.external_configurable(optim.lr_scheduler.CosineAnnealingLR, 'cosine_lr', module='T.optim.lr_scheduler') 41 | config.external_configurable(optim.lr_scheduler.ReduceLROnPlateau, 'plateau_lr', module='T.optim.lr_scheduler') 42 | config.external_configurable(optim.lr_scheduler.CyclicLR, 'cyclic_lr', module='T.optim.lr_scheduler') 43 | config.external_configurable(nnt_optim.lr_scheduler.InverseLR, 'inverse_lr') 44 | config.external_configurable(nnt_optim.lr_scheduler.WarmRestart, 'warm_restart') 45 | 46 | # losses 47 | config.external_configurable(nn.L1Loss, 'l1loss', module='T.nn') 48 | config.external_configurable(nn.MSELoss, 'mseloss', module='T.nn') 49 | config.external_configurable(nn.CrossEntropyLoss, 'celoss', module='T.nn') 50 | config.external_configurable(nn.BCELoss, 'bceloss', module='T.nn') 51 | config.external_configurable(nn.BCEWithLogitsLoss, 'bcelogit_loss', module='T.nn') 52 | config.external_configurable(nn.NLLLoss, 'nllloss', module='T.nn') 53 | config.external_configurable(nn.KLDivLoss, 'kldloss', module='T.nn') 54 | config.external_configurable(nn.SmoothL1Loss, 'huberloss', module='T.nn') 55 | config.external_configurable(nn.CosineEmbeddingLoss, 'cosineembed_loss', module='T.nn') 56 | 57 | # activations 58 | config.external_configurable(nn.ELU, 'elu', module='T.nn') 59 | config.external_configurable(nn.Hardshrink, 'hardshrink', module='T.nn') 60 | config.external_configurable(nn.Hardtanh, 'hardtanh', module='T.nn') 61 | config.external_configurable(nn.LeakyReLU, 'lrelu', module='T.nn') 62 | config.external_configurable(nn.LogSigmoid, 'logsig', module='T.nn') 63 | config.external_configurable(nn.MultiheadAttention, 'multihead_att', module='T.nn') 64 | config.external_configurable(nn.PReLU, 'prelu', module='T.nn') 65 | config.external_configurable(nn.ReLU, 'relu', module='T.nn') 66 | config.external_configurable(nn.RReLU, 'rrelu', module='T.nn') 67 | config.external_configurable(nn.SELU, 'selu', module='T.nn') 68 | config.external_configurable(nn.CELU, 'celu', module='T.nn') 69 | config.external_configurable(nn.Sigmoid, 'sigmoid', module='T.nn') 70 | config.external_configurable(nn.Softplus, 'softplus', module='T.nn') 71 | config.external_configurable(nn.Softshrink, 'softshrink', module='T.nn') 72 | config.external_configurable(nn.Softsign, 'softsign', module='T.nn') 73 | config.external_configurable(nn.Tanh, 'tanh', module='T.nn') 74 | config.external_configurable(nn.Tanhshrink, 'tanhshrink', module='T.nn') 75 | config.external_configurable(nn.Threshold, 'threshold', module='T.nn') 76 | 77 | # constants 78 | config.constant('float16', T.float16) 79 | config.constant('float32', T.float32) 80 | config.constant('float64', T.float64) 81 | config.constant('int8', T.int8) 82 | config.constant('int16', T.int16) 83 | config.constant('int32', T.int32) 84 | config.constant('int64', T.int64) 85 | config.constant('complex32', T.complex32) 86 | config.constant('complex64', T.complex64) 87 | config.constant('complex128', T.complex128) 88 | config.constant('float', T.float) 89 | config.constant('short', T.short) 90 | config.constant('long', T.long) 91 | config.constant('half', T.half) 92 | config.constant('uint8', T.uint8) 93 | config.constant('int', T.int) 94 | 95 | # model zoo 96 | config.external_configurable(zoo.ResNet, 'resnet', module='neuralnet_pytorch.zoo.resnet') 97 | config.external_configurable(zoo.resnet18, 'resnet18', module='neuralnet_pytorch.zoo.resnet') 98 | config.external_configurable(zoo.resnet34, 'resnet34', module='neuralnet_pytorch.zoo.resnet') 99 | config.external_configurable(zoo.resnet50, 'resnet50', module='neuralnet_pytorch.zoo.resnet') 100 | config.external_configurable(zoo.resnet101, 'resnet101', module='neuralnet_pytorch.zoo.resnet') 101 | config.external_configurable(zoo.resnet152, 'resnet152', module='neuralnet_pytorch.zoo.resnet') 102 | config.external_configurable(zoo.resnext50_32x4d, 'resnext50_32x4d', module='neuralnet_pytorch.zoo.resnet') 103 | config.external_configurable(zoo.resnext101_32x8d, 'resnext101_32x8d', module='neuralnet_pytorch.zoo.resnet') 104 | config.external_configurable(zoo.wide_resnet50_2, 'wide_resnet50_2', module='neuralnet_pytorch.zoo.resnet') 105 | config.external_configurable(zoo.wide_resnet101_2, 'wide_resnet101_2', module='neuralnet_pytorch.zoo.resnet') 106 | config.external_configurable(zoo.VGG, 'vgg', module='neuralnet_pytorch.zoo.vgg') 107 | config.external_configurable(zoo.vgg11, 'vgg11', module='neuralnet_pytorch.zoo.vgg') 108 | config.external_configurable(zoo.vgg11_bn, 'vgg11_bn', module='neuralnet_pytorch.zoo.vgg') 109 | config.external_configurable(zoo.vgg13, 'vgg13', module='neuralnet_pytorch.zoo.vgg') 110 | config.external_configurable(zoo.vgg13_bn, 'vgg13_bn', module='neuralnet_pytorch.zoo.vgg') 111 | config.external_configurable(zoo.vgg16, 'vgg16', module='neuralnet_pytorch.zoo.vgg') 112 | config.external_configurable(zoo.vgg16_bn, 'vgg16_bn', module='neuralnet_pytorch.zoo.vgg') 113 | config.external_configurable(zoo.vgg19, 'vgg19', module='neuralnet_pytorch.zoo.vgg') 114 | config.external_configurable(zoo.vgg19_bn, 'vgg19_bn', module='neuralnet_pytorch.zoo.vgg') 115 | -------------------------------------------------------------------------------- /neuralnet_pytorch/utils/cv_utils.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | 3 | from .. import cuda_ext_available 4 | from . import tensor_utils as tutils 5 | 6 | __all__ = ['rgb2gray', 'rgb2ycbcr', 'rgba2rgb', 'ycbcr2rgb', 'pc2vox', 'pc2vox_fast'] 7 | 8 | 9 | def rgb2gray(img: T.Tensor): 10 | """ 11 | Converts a batch of RGB images to gray. 12 | 13 | :param img: 14 | a batch of RGB image tensors. 15 | :return: 16 | a batch of gray images. 17 | """ 18 | 19 | if len(img.shape) != 4: 20 | raise ValueError('Input images must have four dimensions, not %d' % len(img.shape)) 21 | 22 | return (0.299 * img[:, 0] + 0.587 * img[:, 1] + 0.114 * img[:, 2]).unsqueeze(1) 23 | 24 | 25 | def rgb2ycbcr(img: T.Tensor): 26 | """ 27 | Converts a batch of RGB images to YCbCr. 28 | 29 | :param img: 30 | a batch of RGB image tensors. 31 | :return: 32 | a batch of YCbCr images. 33 | """ 34 | 35 | if len(img.shape) != 4: 36 | raise ValueError('Input images must have four dimensions, not %d' % len(img.shape)) 37 | 38 | Y = 0. + .299 * img[:, 0] + .587 * img[:, 1] + .114 * img[:, 2] 39 | Cb = 128. - .169 * img[:, 0] - .331 * img[:, 1] + .5 * img[:, 2] 40 | Cr = 128. + .5 * img[:, 0] - .419 * img[:, 1] - .081 * img[:, 2] 41 | return T.cat((Y.unsqueeze(1), Cb.unsqueeze(1), Cr.unsqueeze(1)), 1) 42 | 43 | 44 | def ycbcr2rgb(img: T.Tensor): 45 | """ 46 | Converts a batch of YCbCr images to RGB. 47 | 48 | :param img: 49 | a batch of YCbCr image tensors. 50 | :return: 51 | a batch of RGB images. 52 | """ 53 | 54 | if len(img.shape) != 4: 55 | raise ValueError('Input images must have four dimensions, not %d' % len(img.shape)) 56 | 57 | R = img[:, 0] + 1.4 * (img[:, 2] - 128.) 58 | G = img[:, 0] - .343 * (img[:, 1] - 128.) - .711 * (img[:, 2] - 128.) 59 | B = img[:, 0] + 1.765 * (img[:, 1] - 128.) 60 | return T.cat((R.unsqueeze(1), G.unsqueeze(1), B.unsqueeze(1)), 1) 61 | 62 | 63 | def rgba2rgb(img: T.Tensor): 64 | """ 65 | Converts a batch of RGBA images to RGB. 66 | 67 | :param img: 68 | an batch of RGBA image tensors. 69 | :return: 70 | a batch of RGB images. 71 | """ 72 | 73 | r = img[..., 0, :, :] 74 | g = img[..., 1, :, :] 75 | b = img[..., 2, :, :] 76 | a = img[..., 3, :, :] 77 | 78 | shape = img.shape[:-3] + (3,) + img.shape[-2:] 79 | out = T.zeros(*shape).to(img.device) 80 | out[..., 0, :, :] = (1 - a) * r + a * r 81 | out[..., 1, :, :] = (1 - a) * g + a * g 82 | out[..., 2, :, :] = (1 - a) * b + a * b 83 | return out 84 | 85 | 86 | def pc2vox(pc: T.Tensor, vox_size=32, sigma=.005, analytical_gauss_norm=True): 87 | """ 88 | Converts a centered point cloud to voxel representation. 89 | 90 | :param pc: 91 | a batch of centered point clouds. 92 | :param vox_size: 93 | resolution of the voxel field. 94 | Default: 32. 95 | :param sigma: 96 | std of the Gaussian blur that determines the area of effect of each point. 97 | :param analytical_gauss_norm: 98 | whether to use a analytically precomputed normalization term. 99 | :return: 100 | the voxel representation of input. 101 | """ 102 | assert pc.ndimension() in (2, 3), 'Point cloud must be a 2D a 3D tensor' 103 | if pc.ndimension() == 2: 104 | pc = pc[None] 105 | 106 | x = pc[..., 0] 107 | y = pc[..., 1] 108 | z = pc[..., 2] 109 | 110 | rng = T.linspace(-1.0, 1.0, vox_size).to(pc.device) 111 | xg, yg, zg = T.meshgrid(rng, rng, rng) # [G,G,G] 112 | 113 | x_big = tutils.shape_padright(x, 3) # [B,N,1,1,1] 114 | y_big = tutils.shape_padright(y, 3) # [B,N,1,1,1] 115 | z_big = tutils.shape_padright(z, 3) # [B,N,1,1,1] 116 | 117 | xg = tutils.shape_padleft(xg, 2) # [1,1,G,G,G] 118 | yg = tutils.shape_padleft(yg, 2) # [1,1,G,G,G] 119 | zg = tutils.shape_padleft(zg, 2) # [1,1,G,G,G] 120 | 121 | # squared distance 122 | sq_distance = (x_big - xg) ** 2. + (y_big - yg) ** 2. + (z_big - zg) ** 2. 123 | 124 | # compute gaussian 125 | func = T.exp(-sq_distance / (2. * sigma ** 2)) # [B,N,G,G,G] 126 | 127 | # normalise gaussian 128 | if analytical_gauss_norm: 129 | # should work with any grid sizes 130 | magic_factor = 1.78984352254 # see estimate_gauss_normaliser 131 | sigma_normalised = sigma * vox_size 132 | normaliser = 1. / (magic_factor * (sigma_normalised ** 3.)) 133 | func *= normaliser 134 | else: 135 | normaliser = T.sum(func, (2, 3, 4), keepdim=True) 136 | func /= normaliser 137 | 138 | summed = T.sum(func, dim=1) # [B,G,G G] 139 | voxels = T.clamp(summed, 0., 1.) 140 | return voxels 141 | 142 | 143 | def pc2vox_fast(pc: T.Tensor, voxel_size=32, grid_size=1., filter_outlier=True, c_code=cuda_ext_available): 144 | """ 145 | A fast conversion from a centered point cloud to voxel representation. 146 | 147 | :param pc: 148 | a batch of centered point clouds. 149 | :param voxel_size: 150 | resolution of the voxel field. 151 | Default: 32. 152 | :param grid_size: 153 | range of the point clouds. 154 | Default: 1. 155 | :param filter_outlier: 156 | whether to filter point outside of `grid_size`. 157 | Default: ``True``. 158 | :param c_code: 159 | whether to use a C++ implementation. 160 | Default: ``True``. 161 | :return: 162 | the voxel representation of input. 163 | """ 164 | 165 | assert pc.ndimension() in (2, 3), 'Point cloud must be a 2D a 3D tensor' 166 | if pc.ndimension() == 2: 167 | pc = pc[None] 168 | 169 | if c_code: 170 | from ..extensions import pc2vox 171 | voxel = pc2vox(pc, voxel_size, grid_size, filter_outlier) 172 | else: 173 | b, n, _ = pc.shape 174 | half_size = grid_size / 2. 175 | valid = (pc >= -half_size) & (pc <= half_size) 176 | valid = T.all(valid, 2) 177 | pc_grid = (pc + half_size) * (voxel_size - 1.) 178 | indices_floor = T.floor(pc_grid) 179 | indices = indices_floor.long() 180 | batch_indices = T.arange(b).to(pc.device) 181 | batch_indices = tutils.shape_padright(batch_indices, 2) 182 | batch_indices = tutils.tile(batch_indices, (1, n, 1)) 183 | indices = T.cat((batch_indices, indices), 2) 184 | indices = T.reshape(indices, (-1, 4)) 185 | 186 | r = pc_grid - indices_floor 187 | rr = (1. - r, r) 188 | if filter_outlier: 189 | valid = valid.flatten() 190 | indices = indices[valid] 191 | 192 | def interpolate_scatter3d(pos): 193 | updates = rr[pos[0]][..., 0] * rr[pos[1]][..., 1] * rr[pos[2]][..., 2] 194 | updates = updates.flatten() 195 | 196 | if filter_outlier: 197 | updates = updates[valid] 198 | 199 | indices_shift = T.tensor([[0] + pos]).to(pc.device) 200 | indices_loc = indices + indices_shift 201 | out_shape = (b,) + (voxel_size,) * 3 202 | voxels = T.zeros(*out_shape).to(pc.device).flatten() 203 | voxels.scatter_add_(0, tutils.ravel_index(indices_loc.t(), out_shape), updates) 204 | return voxels.view(*out_shape) 205 | 206 | voxel = [interpolate_scatter3d([k, j, i]) for k in range(2) for j in range(2) for i in range(2)] 207 | voxel = sum(voxel) 208 | voxel = T.clamp(voxel, 0., 1.) 209 | return voxel 210 | -------------------------------------------------------------------------------- /tests/function_test.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | import numpy as np 3 | from torch import testing 4 | import pytest 5 | 6 | import neuralnet_pytorch as nnt 7 | from neuralnet_pytorch import cuda_available 8 | 9 | dev = ('cpu', 'cuda') if cuda_available else ('cpu',) 10 | 11 | 12 | @pytest.mark.parametrize('device', dev) 13 | def test_ravel_index(device): 14 | shape = (2, 4, 5, 3) 15 | a = T.arange(np.prod(shape)).reshape(*shape).to(device) 16 | 17 | indices = [[1, 0, 1, 1, 0], [1, 3, 3, 2, 1], [1, 1, 4, 0, 3], [1, 2, 2, 2, 0]] 18 | linear_indices = nnt.utils.ravel_index(T.tensor(indices), shape) 19 | testing.assert_allclose(linear_indices.type_as(a), a[indices]) 20 | 21 | 22 | @pytest.mark.parametrize('device', dev) 23 | def test_shape_pad(device): 24 | shape = (10, 10) 25 | a = T.rand(*shape).to(device) 26 | 27 | padded = nnt.utils.shape_padleft(a) 28 | expected = a.unsqueeze(0) 29 | testing.assert_allclose(padded, expected) 30 | testing.assert_allclose(padded.shape, expected.shape) 31 | 32 | padded = nnt.utils.shape_padleft(a, 2) 33 | expected = a.unsqueeze(0).unsqueeze(0) 34 | testing.assert_allclose(padded, expected) 35 | testing.assert_allclose(padded.shape, expected.shape) 36 | 37 | padded = nnt.utils.shape_padleft(a, 5) 38 | expected = a.unsqueeze(0).unsqueeze(0).unsqueeze(0).unsqueeze(0).unsqueeze(0) 39 | testing.assert_allclose(padded, expected) 40 | testing.assert_allclose(padded.shape, expected.shape) 41 | 42 | padded = nnt.utils.shape_padright(a) 43 | expected = a.unsqueeze(-1) 44 | testing.assert_allclose(padded, expected) 45 | testing.assert_allclose(padded.shape, expected.shape) 46 | 47 | padded = nnt.utils.shape_padright(a, 2) 48 | expected = a.unsqueeze(-1).unsqueeze(-1) 49 | testing.assert_allclose(padded, expected) 50 | testing.assert_allclose(padded.shape, expected.shape) 51 | 52 | padded = nnt.utils.shape_padright(a, 5) 53 | expected = a.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) 54 | testing.assert_allclose(padded, expected) 55 | testing.assert_allclose(padded.shape, expected.shape) 56 | 57 | 58 | @pytest.mark.parametrize('device', dev) 59 | def test_dimshuffle(device): 60 | shape = (64, 512) 61 | a = T.rand(*shape).to(device) 62 | 63 | dimshuffled = nnt.utils.dimshuffle(a, (0, 1, 'x', 'x')) 64 | expected = a.unsqueeze(2).unsqueeze(2) 65 | testing.assert_allclose(dimshuffled, expected) 66 | testing.assert_allclose(dimshuffled.shape, expected.shape) 67 | 68 | dimshuffled = nnt.utils.dimshuffle(a, (1, 0, 'x', 'x')) 69 | expected = a.permute(1, 0).unsqueeze(2).unsqueeze(2) 70 | testing.assert_allclose(dimshuffled, expected) 71 | testing.assert_allclose(dimshuffled.shape, expected.shape) 72 | 73 | dimshuffled = nnt.utils.dimshuffle(a, (0, 'x', 1, 'x')) 74 | expected = a.unsqueeze(2).permute(0, 2, 1).unsqueeze(3) 75 | testing.assert_allclose(dimshuffled, expected) 76 | testing.assert_allclose(dimshuffled.shape, expected.shape) 77 | 78 | dimshuffled = nnt.utils.dimshuffle(a, (1, 'x', 'x', 0)) 79 | expected = a.permute(1, 0).unsqueeze(1).unsqueeze(1) 80 | testing.assert_allclose(dimshuffled, expected) 81 | testing.assert_allclose(dimshuffled.shape, expected.shape) 82 | 83 | dimshuffled = nnt.utils.dimshuffle(a, (1, 'x', 0, 'x', 'x')) 84 | expected = a.permute(1, 0).unsqueeze(1).unsqueeze(3).unsqueeze(3) 85 | testing.assert_allclose(dimshuffled, expected) 86 | testing.assert_allclose(dimshuffled.shape, expected.shape) 87 | 88 | 89 | @pytest.mark.parametrize('device', dev) 90 | def test_flatten(device): 91 | shape = (10, 4, 2, 3, 6) 92 | a = T.rand(*shape).to(device) 93 | 94 | flatten = nnt.Flatten(2, input_shape=(None,) + shape[1:]) 95 | expected = T.flatten(a, 2) 96 | testing.assert_allclose(flatten(a), expected) 97 | testing.assert_allclose((shape[0],) + flatten.output_shape[1:], expected.shape) 98 | 99 | flatten = nnt.Flatten(4, input_shape=(None,) + shape[1:]) 100 | expected = T.flatten(a, 4) 101 | testing.assert_allclose(flatten(a), expected) 102 | testing.assert_allclose((shape[0],) + flatten.output_shape[1:], expected.shape) 103 | 104 | flatten = nnt.Flatten(input_shape=shape) 105 | expected = T.flatten(a) 106 | testing.assert_allclose(flatten(a), expected) 107 | testing.assert_allclose(flatten.output_shape, expected.shape) 108 | 109 | flatten = nnt.Flatten(1, 3, input_shape=shape) 110 | expected = T.flatten(a, 1, 3) 111 | testing.assert_allclose(flatten(a), expected) 112 | testing.assert_allclose(flatten.output_shape, expected.shape) 113 | 114 | 115 | @pytest.mark.parametrize('device', dev) 116 | def test_reshape(device): 117 | shape = (10, 3, 9, 9) 118 | a = T.rand(*shape).to(device) 119 | 120 | newshape = (-1, 9, 9) 121 | reshape = nnt.Reshape(newshape, input_shape=(None,) + shape[1:]) 122 | expected = T.reshape(a, newshape) 123 | testing.assert_allclose(reshape(a), expected) 124 | testing.assert_allclose(reshape.output_shape[1:], expected.shape[1:]) 125 | 126 | newshape = (10, -1, 9) 127 | reshape = nnt.Reshape(newshape, input_shape=shape) 128 | expected = T.reshape(a, newshape) 129 | testing.assert_allclose(reshape(a), expected) 130 | testing.assert_allclose(reshape.output_shape, expected.shape) 131 | 132 | newshape = (9, 9, -1) 133 | reshape = nnt.Reshape(newshape, input_shape=shape) 134 | expected = T.reshape(a, newshape) 135 | testing.assert_allclose(reshape(a), expected) 136 | testing.assert_allclose(reshape.output_shape[1:], expected.shape[1:]) 137 | 138 | 139 | @pytest.mark.parametrize('device', dev) 140 | def test_batch_pairwise_distance(device): 141 | xyz = T.rand(1, 4, 3).to(device) 142 | actual = nnt.utils.batch_pairwise_sqdist(xyz, xyz, c_code=False) 143 | testing.assert_allclose(T.diag(actual[0]), T.zeros(actual.shape[1]).to(device)) 144 | 145 | if dev != 'cuda' or not nnt.cuda_ext_available: 146 | pytest.skip('Requires CUDA extension to be installed') 147 | 148 | xyz1 = T.rand(10, 4000, 3).to(device).requires_grad_(True) 149 | xyz2 = T.rand(10, 5000, 3).to(device) 150 | expected = nnt.utils.batch_pairwise_sqdist(xyz1, xyz2, c_code=False) 151 | actual = nnt.utils.batch_pairwise_sqdist(xyz1, xyz2, c_code=True) 152 | testing.assert_allclose(actual, expected) 153 | 154 | expected_cost = T.sum(expected) 155 | expected_cost.backward() 156 | expected_grad = xyz1.grad 157 | xyz1.grad.zero_() 158 | 159 | actual_cost = T.sum(actual) 160 | actual_cost.backward() 161 | actual_grad = xyz1.grad 162 | testing.assert_allclose(actual_grad, expected_grad) 163 | 164 | for _ in range(10): 165 | t1 = nnt.utils.time_cuda_module(nnt.utils.batch_pairwise_sqdist, xyz1, xyz2, c_code=False) 166 | t2 = nnt.utils.time_cuda_module(nnt.utils.batch_pairwise_sqdist, xyz1, xyz2, c_code=True) 167 | print('pt: %f, cpp: %f' % (t1, t2)) 168 | 169 | 170 | @pytest.mark.skipif(not nnt.cuda_ext_available, reason='Requires CUDA extension to be installed') 171 | @pytest.mark.parametrize('device', dev) 172 | def test_pointcloud_to_voxel(device): 173 | xyz = T.rand(10, 4000, 3).to(device).requires_grad_(True) 174 | pc = xyz * 2. - 1. 175 | expected = nnt.utils.pc2vox_fast(pc, c_code=False) 176 | actual = nnt.utils.pc2vox_fast(pc, c_code=True) 177 | testing.assert_allclose(actual, expected) 178 | 179 | expected_cost = T.sum(expected) 180 | expected_cost.backward(retain_graph=True) 181 | expected_grad = xyz.grad 182 | xyz.grad.zero_() 183 | 184 | actual_cost = T.sum(actual) 185 | actual_cost.backward() 186 | actual_grad = xyz.grad 187 | testing.assert_allclose(actual_grad, expected_grad) 188 | 189 | for _ in range(10): 190 | t1 = nnt.utils.time_cuda_module(nnt.utils.pc2vox_fast, pc, c_code=False) 191 | t2 = nnt.utils.time_cuda_module(nnt.utils.pc2vox_fast, pc) 192 | print('pt: %f, cpp: %f' % (t1, t2)) 193 | -------------------------------------------------------------------------------- /neuralnet_pytorch/zoo/vgg.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torchvision.models.vgg import model_urls 3 | from torchvision.models.utils import load_state_dict_from_url 4 | from ..layers import Sequential, Conv2d, ConvNormAct, FC, MaxPool2d, wrapper 5 | from ..utils import batch_set_tensor 6 | 7 | __all__ = [ 8 | 'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 9 | 'vgg19_bn', 'vgg19', 10 | ] 11 | 12 | 13 | class VGG(Sequential): 14 | 15 | def __init__(self, features, num_classes=1000, default_init=True): 16 | super(VGG, self).__init__() 17 | self.features = features 18 | self.avgpool = wrapper(output_size=(7, 7))(nn.AdaptiveAvgPool2d)() 19 | self.classifier = Sequential( 20 | FC(512 * 7 * 7, 4096, flatten=True, activation='relu'), 21 | wrapper()(nn.Dropout)(), 22 | FC(4096, 4096, activation='relu'), 23 | wrapper()(nn.Dropout)(), 24 | FC(4096, num_classes), 25 | ) 26 | if default_init: 27 | self._initialize_weights() 28 | 29 | def _initialize_weights(self): 30 | for m in self.modules(): 31 | if isinstance(m, nn.Conv2d): 32 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 33 | if m.bias is not None: 34 | nn.init.constant_(m.bias, 0) 35 | elif isinstance(m, nn.BatchNorm2d): 36 | nn.init.constant_(m.weight, 1) 37 | nn.init.constant_(m.bias, 0) 38 | elif isinstance(m, nn.Linear): 39 | nn.init.normal_(m.weight, 0, 0.01) 40 | nn.init.constant_(m.bias, 0) 41 | 42 | 43 | def make_layers(cfg, batch_norm=False): 44 | layers = [] 45 | in_channels = 3 46 | for v in cfg: 47 | if v == 'M': 48 | layers += [MaxPool2d(kernel_size=2, stride=2)] 49 | else: 50 | conv = ConvNormAct if batch_norm else Conv2d 51 | layers += [conv(in_channels, v, kernel_size=3, padding=1, activation='relu')] 52 | in_channels = v 53 | return Sequential(*layers) 54 | 55 | 56 | cfgs = { 57 | 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 58 | 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 59 | 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 60 | 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], 61 | } 62 | 63 | 64 | def _vgg(arch, cfg, batch_norm, pretrained, progress, **kwargs): 65 | if pretrained: 66 | kwargs['default_init'] = False 67 | 68 | model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs) 69 | if pretrained: 70 | state_dict = load_state_dict_from_url(model_urls[arch], 71 | progress=progress) 72 | try: 73 | batch_set_tensor(model.state_dict().values(), state_dict.values()) 74 | except (RuntimeError, ValueError): 75 | state_dict_iter = iter(state_dict.items()) 76 | for k, v in model.state_dict().items(): 77 | if 'num_batches_tracked' not in k: 78 | k_t, v_t = next(state_dict_iter) 79 | param_name = k.split('.')[-1] 80 | value_name = k_t.split('.') 81 | if param_name != value_name[-1]: 82 | value_name[-1] = param_name 83 | value_name = '.'.join(value_name) 84 | v.data.copy_(state_dict[value_name].data) 85 | else: 86 | v.data.copy_(v_t.data) 87 | return model 88 | 89 | 90 | def vgg11(pretrained=False, progress=True, **kwargs): 91 | """ 92 | VGG 11-layer model (configuration "A") from 93 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ 94 | 95 | :param pretrained: 96 | If True, returns a model pre-trained on ImageNet. 97 | :param progress: 98 | If True, displays a progress bar of the download to stderr. 99 | """ 100 | return _vgg('vgg11', 'A', False, pretrained, progress, **kwargs) 101 | 102 | 103 | def vgg11_bn(pretrained=False, progress=True, **kwargs): 104 | """ 105 | VGG 11-layer model (configuration "A") with batch normalization 106 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ 107 | 108 | :param pretrained: 109 | If True, returns a model pre-trained on ImageNet. 110 | :param progress: 111 | If True, displays a progress bar of the download to stderr. 112 | """ 113 | return _vgg('vgg11_bn', 'A', True, pretrained, progress, **kwargs) 114 | 115 | 116 | def vgg13(pretrained=False, progress=True, **kwargs): 117 | """ 118 | VGG 13-layer model (configuration "B") 119 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ 120 | 121 | :param pretrained: 122 | If True, returns a model pre-trained on ImageNet. 123 | :param progress: 124 | If True, displays a progress bar of the download to stderr. 125 | """ 126 | return _vgg('vgg13', 'B', False, pretrained, progress, **kwargs) 127 | 128 | 129 | def vgg13_bn(pretrained=False, progress=True, **kwargs): 130 | """ 131 | VGG 13-layer model (configuration "B") with batch normalization 132 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ 133 | 134 | :param pretrained: 135 | If True, returns a model pre-trained on ImageNet. 136 | :param progress: 137 | If True, displays a progress bar of the download to stderr. 138 | """ 139 | return _vgg('vgg13_bn', 'B', True, pretrained, progress, **kwargs) 140 | 141 | 142 | def vgg16(pretrained=False, progress=True, **kwargs): 143 | """ 144 | VGG 16-layer model (configuration "D") 145 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ 146 | 147 | :param pretrained: 148 | If True, returns a model pre-trained on ImageNet. 149 | :param progress: 150 | If True, displays a progress bar of the download to stderr. 151 | """ 152 | return _vgg('vgg16', 'D', False, pretrained, progress, **kwargs) 153 | 154 | 155 | def vgg16_bn(pretrained=False, progress=True, **kwargs): 156 | """VGG 16-layer model (configuration "D") with batch normalization 157 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ 158 | 159 | :param pretrained: 160 | If True, returns a model pre-trained on ImageNet. 161 | :param progress: 162 | If True, displays a progress bar of the download to stderr. 163 | """ 164 | return _vgg('vgg16_bn', 'D', True, pretrained, progress, **kwargs) 165 | 166 | 167 | def vgg19(pretrained=False, progress=True, **kwargs): 168 | """ 169 | VGG 19-layer model (configuration "E") 170 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ 171 | 172 | :param pretrained: 173 | If True, returns a model pre-trained on ImageNet. 174 | :param progress: 175 | If True, displays a progress bar of the download to stderr. 176 | """ 177 | return _vgg('vgg19', 'E', False, pretrained, progress, **kwargs) 178 | 179 | 180 | def vgg19_bn(pretrained=False, progress=True, **kwargs): 181 | """ 182 | VGG 19-layer model (configuration 'E') with batch normalization 183 | `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ 184 | 185 | :param pretrained: 186 | If True, returns a model pre-trained on ImageNet. 187 | :param progress: 188 | If True, displays a progress bar of the download to stderr. 189 | """ 190 | return _vgg('vgg19_bn', 'E', True, pretrained, progress, **kwargs) 191 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/csrc/chamfer_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | template 10 | __global__ void 11 | NmDistanceKernel(int b, int n, const scalar_t* __restrict__ xyz, int m, 12 | const scalar_t* __restrict__ xyz2, 13 | scalar_t* __restrict__ result, scalar_t* __restrict__ result_i) 14 | { 15 | const int batch = 512; 16 | __shared__ scalar_t buf[batch * 3]; 17 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 18 | for (int k2 = 0; k2 < m; k2 += batch) { 19 | int end_k = min(m, k2 + batch) - k2; 20 | for (int j = threadIdx.x; j < end_k * 3; j += blockDim.x) { 21 | buf[j] = xyz2[(i * m + k2) * 3 + j]; 22 | } 23 | __syncthreads(); 24 | for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; 25 | j += blockDim.x * gridDim.y) { 26 | scalar_t x1 = xyz[(i * n + j) * 3 + 0]; 27 | scalar_t y1 = xyz[(i * n + j) * 3 + 1]; 28 | scalar_t z1 = xyz[(i * n + j) * 3 + 2]; 29 | int best_i = 0; 30 | scalar_t best = 0; 31 | int end_ka = end_k - (end_k & 3); 32 | if (end_ka == batch) { 33 | for (int k = 0; k < batch; k += 4) { 34 | { 35 | scalar_t x2 = buf[k * 3 + 0] - x1; 36 | scalar_t y2 = buf[k * 3 + 1] - y1; 37 | scalar_t z2 = buf[k * 3 + 2] - z1; 38 | scalar_t d = x2 * x2 + y2 * y2 + z2 * z2; 39 | if (k == 0 || d < best) { 40 | best = d; 41 | best_i = k + k2; 42 | } 43 | } 44 | { 45 | scalar_t x2 = buf[k * 3 + 3] - x1; 46 | scalar_t y2 = buf[k * 3 + 4] - y1; 47 | scalar_t z2 = buf[k * 3 + 5] - z1; 48 | scalar_t d = x2 * x2 + y2 * y2 + z2 * z2; 49 | if (d < best) { 50 | best = d; 51 | best_i = k + k2 + 1; 52 | } 53 | } 54 | { 55 | scalar_t x2 = buf[k * 3 + 6] - x1; 56 | scalar_t y2 = buf[k * 3 + 7] - y1; 57 | scalar_t z2 = buf[k * 3 + 8] - z1; 58 | scalar_t d = x2 * x2 + y2 * y2 + z2 * z2; 59 | if (d < best) { 60 | best = d; 61 | best_i = k + k2 + 2; 62 | } 63 | } 64 | { 65 | scalar_t x2 = buf[k * 3 + 9] - x1; 66 | scalar_t y2 = buf[k * 3 + 10] - y1; 67 | scalar_t z2 = buf[k * 3 + 11] - z1; 68 | scalar_t d = x2 * x2 + y2 * y2 + z2 * z2; 69 | if (d < best) { 70 | best = d; 71 | best_i = k + k2 + 3; 72 | } 73 | } 74 | } 75 | } else { 76 | for (int k = 0; k < end_ka; k += 4) { 77 | { 78 | scalar_t x2 = buf[k * 3 + 0] - x1; 79 | scalar_t y2 = buf[k * 3 + 1] - y1; 80 | scalar_t z2 = buf[k * 3 + 2] - z1; 81 | scalar_t d = x2 * x2 + y2 * y2 + z2 * z2; 82 | if (k == 0 || d < best) { 83 | best = d; 84 | best_i = k + k2; 85 | } 86 | } 87 | { 88 | scalar_t x2 = buf[k * 3 + 3] - x1; 89 | scalar_t y2 = buf[k * 3 + 4] - y1; 90 | scalar_t z2 = buf[k * 3 + 5] - z1; 91 | scalar_t d = x2 * x2 + y2 * y2 + z2 * z2; 92 | if (d < best) { 93 | best = d; 94 | best_i = k + k2 + 1; 95 | } 96 | } 97 | { 98 | scalar_t x2 = buf[k * 3 + 6] - x1; 99 | scalar_t y2 = buf[k * 3 + 7] - y1; 100 | scalar_t z2 = buf[k * 3 + 8] - z1; 101 | scalar_t d = x2 * x2 + y2 * y2 + z2 * z2; 102 | if (d < best) { 103 | best = d; 104 | best_i = k + k2 + 2; 105 | } 106 | } 107 | { 108 | scalar_t x2 = buf[k * 3 + 9] - x1; 109 | scalar_t y2 = buf[k * 3 + 10] - y1; 110 | scalar_t z2 = buf[k * 3 + 11] - z1; 111 | scalar_t d = x2 * x2 + y2 * y2 + z2 * z2; 112 | if (d < best) { 113 | best = d; 114 | best_i = k + k2 + 3; 115 | } 116 | } 117 | } 118 | } 119 | for (int k = end_ka; k < end_k; k++) { 120 | scalar_t x2 = buf[k * 3 + 0] - x1; 121 | scalar_t y2 = buf[k * 3 + 1] - y1; 122 | scalar_t z2 = buf[k * 3 + 2] - z1; 123 | scalar_t d = x2 * x2 + y2 * y2 + z2 * z2; 124 | if (k == 0 || d < best) { 125 | best = d; 126 | best_i = k + k2; 127 | } 128 | } 129 | if (k2 == 0 || result[(i * n + j)] > best) { 130 | result[(i * n + j)] = best; 131 | result_i[(i * n + j)] = scalar_t(best_i); 132 | } 133 | } 134 | __syncthreads(); 135 | } 136 | } 137 | } 138 | 139 | std::vector 140 | chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2) 141 | { 142 | cudaSetDevice((int)xyz1.device().index()); 143 | const auto batch_size = xyz1.size(0); 144 | const auto n = xyz1.size(1); // num_points point cloud A 145 | const auto m = xyz2.size(1); // num_points point cloud B 146 | 147 | auto dist1 = torch::zeros({ batch_size, n }, xyz1.type()); 148 | auto dist2 = torch::zeros({ batch_size, m }, xyz1.type()); 149 | auto idx1 = torch::zeros({ batch_size, n }, xyz1.type()); 150 | auto idx2 = torch::zeros({ batch_size, m }, xyz1.type()); 151 | 152 | AT_DISPATCH_FLOATING_TYPES_AND( 153 | at::ScalarType::Half, xyz1.scalar_type(), "chamfer_cuda_forward", ([&] { 154 | NmDistanceKernel<<>>( 155 | batch_size, n, xyz1.data(), m, xyz2.data(), 156 | dist1.data(), idx1.data()); 157 | NmDistanceKernel<<>>( 158 | batch_size, m, xyz2.data(), n, xyz1.data(), 159 | dist2.data(), idx2.data()); 160 | })); 161 | THCudaCheck(cudaGetLastError()); 162 | 163 | return { dist1, dist2, idx1, idx2 }; 164 | } 165 | 166 | template 167 | __global__ void 168 | NmDistanceGradKernel(int b, int n, const scalar_t* __restrict__ xyz1, int m, 169 | const scalar_t* __restrict__ xyz2, 170 | const scalar_t* __restrict__ grad_dist1, 171 | const scalar_t* __restrict__ idx, 172 | scalar_t* __restrict__ grad_xyz1, 173 | scalar_t* __restrict__ grad_xyz2) 174 | { 175 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 176 | for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; 177 | j += blockDim.x * gridDim.y) { 178 | scalar_t x1 = xyz1[(i * n + j) * 3 + 0]; 179 | scalar_t y1 = xyz1[(i * n + j) * 3 + 1]; 180 | scalar_t z1 = xyz1[(i * n + j) * 3 + 2]; 181 | int j2 = (int)idx[i * n + j]; 182 | scalar_t x2 = xyz2[(i * m + j2) * 3 + 0]; 183 | scalar_t y2 = xyz2[(i * m + j2) * 3 + 1]; 184 | scalar_t z2 = xyz2[(i * m + j2) * 3 + 2]; 185 | scalar_t g = grad_dist1[i * n + j] * 2; 186 | atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 0]), g * (x1 - x2)); 187 | atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 1]), g * (y1 - y2)); 188 | atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 2]), g * (z1 - z2)); 189 | atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 0]), -(g * (x1 - x2))); 190 | atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 1]), -(g * (y1 - y2))); 191 | atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 2]), -(g * (z1 - z2))); 192 | } 193 | } 194 | } 195 | 196 | std::vector 197 | chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor graddist1, 198 | at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) 199 | { 200 | const auto batch_size = xyz1.size(0); 201 | const auto n = xyz1.size(1); // num_points point cloud A 202 | const auto m = xyz2.size(1); // num_points point cloud B 203 | 204 | auto gradxyz1 = torch::zeros_like(xyz1); 205 | auto gradxyz2 = torch::zeros_like(xyz2); 206 | 207 | AT_DISPATCH_FLOATING_TYPES_AND( 208 | at::ScalarType::Half, xyz1.scalar_type(), "chamfer_backward_cuda", ([&] { 209 | NmDistanceGradKernel<<>>( 210 | batch_size, n, xyz1.data(), m, xyz2.data(), 211 | graddist1.data(), idx1.data(), 212 | gradxyz1.data(), gradxyz2.data()); 213 | NmDistanceGradKernel<<>>( 214 | batch_size, m, xyz2.data(), n, xyz1.data(), 215 | graddist2.data(), idx2.data(), 216 | gradxyz2.data(), gradxyz1.data()); 217 | })); 218 | THCudaCheck(cudaGetLastError()); 219 | 220 | return { gradxyz1, gradxyz2 }; 221 | } 222 | -------------------------------------------------------------------------------- /neuralnet_pytorch/metrics.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | import torch.nn.functional as F 3 | import numpy as np 4 | import torch.nn as nn 5 | 6 | from . import layers 7 | from . import utils 8 | from . import cuda_ext_available 9 | 10 | __all__ = ['huber_loss', 'first_derivative_loss', 'lp_loss', 'ssim', 'psnr', 'chamfer_loss', 'emd_loss', 'tv_reg', 11 | 'spectral_norm'] 12 | 13 | 14 | def huber_loss(x, y, reduce='mean'): 15 | """ 16 | An alias for :func:`torch.nn.functional.smooth_l1_loss`. 17 | """ 18 | 19 | return F.smooth_l1_loss(x, y, reduce=reduce) 20 | 21 | 22 | def first_derivative_loss(x, y, p=2): 23 | """ 24 | Calculates lp loss between the first derivatives of the inputs. 25 | 26 | :param x: 27 | a :class:`torch.Tensor`. 28 | :param y: 29 | a :class:`torch.Tensor` of the same shape as x. 30 | :param p: 31 | order of the norm. 32 | :return: 33 | the scalar loss between the first derivatives of the inputs. 34 | """ 35 | 36 | if x.ndimension() != 4 and y.ndimension() != 4: 37 | raise TypeError('y and y_pred should have four dimensions') 38 | 39 | kern_x = T.from_numpy(np.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype='float32')).requires_grad_(False) 40 | kern_x = T.flip(kern_x.expand(y.shape[1], y.shape[1], 3, 3), (0, 1)).to(x.device) 41 | 42 | kern_y = T.from_numpy(np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype='float32')).requires_grad_(False) 43 | kern_y = T.flip(kern_y.expand(y.shape[1], y.shape[1], 3, 3), (0, 1)).to(x.device) 44 | 45 | x_grad_x = F.conv2d(x, kern_x, padding=1) 46 | x_grad_y = F.conv2d(x, kern_y, padding=1) 47 | x_grad = T.sqrt(x_grad_x ** 2 + x_grad_y ** 2 + 1e-10) 48 | 49 | y_grad_x = F.conv2d(y, kern_x, padding=1) 50 | y_grad_y = F.conv2d(y, kern_y, padding=1) 51 | y_grad = T.sqrt(y_grad_x ** 2 + y_grad_y ** 2 + 1e-10) 52 | return lp_loss(x_grad, y_grad, p) 53 | 54 | 55 | def lp_loss(x, y, p=2, reduction='mean'): 56 | """ 57 | Calculates p-norm of (x - y). 58 | 59 | :param x: 60 | a :class:`torch.Tensor`. 61 | :param y: 62 | a :class:`torch.Tensor` of the same shape as x. 63 | :param p: 64 | order of the norm. 65 | :param reduction: 66 | ``'mean'`` or ``'sum'``. 67 | :return: 68 | the p-norm of (x - y). 69 | """ 70 | assert reduction in ('sum', 'mean'), 'Unknown choice of reduction' 71 | 72 | if y.ndimension() != x.ndimension(): 73 | raise TypeError('y should have the same shape as y_pred', ('y', y.data.type(), 'y_pred', x.data.type())) 74 | 75 | if p == 1: 76 | return F.l1_loss(x, y, reduction=reduction) 77 | elif p == 2: 78 | return F.mse_loss(x, y, reduction=reduction) 79 | else: 80 | reduce = T.mean if reduction == 'mean' else T.sum 81 | return reduce(T.abs(x - y) ** p) 82 | 83 | 84 | def chamfer_loss(xyz1, xyz2, reduce='mean', c_code=cuda_ext_available): 85 | """ 86 | Calculates the Chamfer distance between two batches of point clouds. 87 | The Pytorch code is adapted from DenseLidarNet_. 88 | The CUDA code is adapted from AtlasNet_. 89 | 90 | .. _DenseLidarNet: https://github.com/345ishaan/DenseLidarNet/blob/master/code/chamfer_loss.py 91 | .. _AtlasNet: https://github.com/ThibaultGROUEIX/AtlasNet/tree/master/extension 92 | 93 | :param xyz1: 94 | a point cloud of shape ``(b, n1, k)`` or ``(n1, k)``. 95 | :param xyz2: 96 | a point cloud of shape (b, n2, k) or (n2, k). 97 | :param reduce: 98 | ``'mean'`` or ``'sum'``. Default: ``'mean'``. 99 | :param c_code: 100 | whether to use CUDA implementation. 101 | This version is much more memory-friendly and slightly faster. 102 | :return: 103 | the Chamfer distance between the inputs. 104 | """ 105 | assert len(xyz1.shape) in (2, 3) and len(xyz2.shape) in (2, 3), 'Unknown shape of tensors' 106 | 107 | if xyz1.dim() == 2: 108 | xyz1 = xyz1.unsqueeze(0) 109 | 110 | if xyz2.dim() == 2: 111 | xyz2 = xyz2.unsqueeze(0) 112 | 113 | assert reduce in ('mean', 'sum'), 'Unknown reduce method' 114 | reduce = T.sum if reduce == 'sum' else T.mean 115 | 116 | if c_code: 117 | from .extensions import chamfer_distance 118 | dist1, dist2 = chamfer_distance(xyz1, xyz2) 119 | else: 120 | P = utils.batch_pairwise_sqdist(xyz1, xyz2, c_code=c_code) 121 | dist2, _ = T.min(P, 1) 122 | dist1, _ = T.min(P, 2) 123 | loss_2 = reduce(dist2) 124 | loss_1 = reduce(dist1) 125 | return loss_1 + loss_2 126 | 127 | 128 | def emd_loss(xyz1, xyz2, reduce='mean', sinkhorn=False): 129 | """ 130 | Calculates the Earth Mover Distance (or Wasserstein metric) between two sets 131 | of points. 132 | 133 | :param xyz1: 134 | a point cloud of shape ``(b, n1, k)`` or ``(n1, k)``. 135 | :param xyz2: 136 | a point cloud of shape (b, n2, k) or (n2, k). 137 | :param reduce: 138 | ``'mean'`` or ``'sum'``. Default: ``'mean'``. 139 | :param sinkhorn: 140 | whether to use the Sinkhorn approximation of the Wasserstein distance. 141 | ``False`` will fall back to a CUDA implementation, which is only available 142 | if the CUDA-extended neuralnet-pytorch is installed. 143 | Default: ``True``. 144 | :return: 145 | the EMD between the inputs. 146 | """ 147 | 148 | assert reduce in ('mean', 'sum'), 'Reduce method should be mean or sum' 149 | if sinkhorn: 150 | import geomloss 151 | return geomloss.SamplesLoss()(xyz1, xyz2) 152 | else: 153 | from .extensions import earth_mover_distance as emd 154 | emd_dist = (emd(xyz1, xyz2) + emd(xyz2, xyz1)) / 2. 155 | return T.mean(emd_dist) if reduce == 'mean' else T.sum(emd_dist) 156 | 157 | 158 | def _fspecial_gauss(size, sigma): 159 | x, y = np.mgrid[-size // 2 + 1:size // 2 + 1, -size // 2 + 1:size // 2 + 1] 160 | g = np.exp(-((x ** 2 + y ** 2) / (2.0 * sigma ** 2))) 161 | return g / np.sum(g) 162 | 163 | 164 | def ssim(img1, img2, max_val=1., filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03, cs_map=False): 165 | """ 166 | Returns the Structural Similarity Map between `img1` and `img2`. 167 | This function attempts to match the functionality of ssim_index_new.m by 168 | Zhou Wang: http://www.cns.nyu.edu/~lcv/ssim/msssim.zip 169 | 170 | :param img1: 171 | a 4D :class:`torch.Tensor`. 172 | :param img2: 173 | a 4D :class:`torch.Tensor` of the same shape as `img1`. 174 | :param max_val: 175 | the dynamic range of the images (i.e., the difference between the 176 | maximum the and minimum allowed values). 177 | :param filter_size: 178 | size of blur kernel to use (will be reduced for small images). 179 | :param filter_sigma: 180 | standard deviation for Gaussian blur kernel (will be reduced 181 | for small images). 182 | :param k1: 183 | constant used to maintain stability in the SSIM calculation (0.01 in 184 | the original paper). 185 | :param k2: 186 | constant used to maintain stability in the SSIM calculation (0.03 in 187 | the original paper). 188 | :return: 189 | pair containing the mean SSIM and contrast sensitivity between `img1` and `img2`. 190 | :raise: 191 | RuntimeError: If input images don't have the same shape or don't have four 192 | dimensions: [batch_size, height, width, depth]. 193 | """ 194 | 195 | if img1.ndimension() != 4: 196 | raise RuntimeError('Input images must have four dimensions, not %d', img1.ndimension()) 197 | 198 | _, _, height, width = img1.shape 199 | 200 | # Filter size can't be larger than height or width of images. 201 | size = min((filter_size, height, width)) 202 | 203 | # Scale down sigma if a smaller filter size is used. 204 | sigma = (size * filter_sigma / filter_size) if filter_size else 1. 205 | 206 | if filter_size: 207 | window = T.flip(T.tensor(_fspecial_gauss(size, sigma)), (0, 1)).view(1, 1, size, size)\ 208 | .requires_grad_(False).to(device=img1.device, dtype=img1.dtype) 209 | 210 | mu1 = F.conv2d(img1, window) 211 | mu2 = F.conv2d(img2, window) 212 | sigma11 = F.conv2d(img1 * img1, window) 213 | sigma22 = F.conv2d(img2 * img2, window) 214 | sigma12 = F.conv2d(img1 * img2, window) 215 | else: 216 | # Empty blur kernel so no need to convolve. 217 | mu1, mu2 = img1, img2 218 | sigma11 = img1 * img1 219 | sigma22 = img2 * img2 220 | sigma12 = img1 * img2 221 | 222 | mu1 = mu1 * mu1 223 | mu2 = mu2 * mu2 224 | mu12 = mu1 * mu2 225 | sigma11 -= mu1 226 | sigma22 -= mu2 227 | sigma12 -= mu12 228 | 229 | # Calculate intermediate values used by both ssim and cs_map. 230 | c1 = (k1 * max_val) ** 2 231 | c2 = (k2 * max_val) ** 2 232 | v1 = 2.0 * sigma12 + c2 233 | v2 = sigma11 + sigma22 + c2 234 | ssim = T.mean((((2.0 * mu12 + c1) * v1) / ((mu1 + mu2 + c1) * v2 + 1e-10))) 235 | output = ssim if not cs_map else (ssim, T.mean(v1 / v2)) 236 | return output 237 | 238 | 239 | def psnr(x, y): 240 | """ 241 | Peak-signal-to-noise ratio for [0,1] images. 242 | 243 | :param x: 244 | a :class:`torch.Tensor`. 245 | :param y: 246 | a :class:`torch.Tensor` of the same shape as `x`. 247 | """ 248 | 249 | return -10 * T.log(T.mean((y - x) ** 2)) / np.log(10.) 250 | 251 | 252 | def tv_reg(y): 253 | """ 254 | Total variation regularization. 255 | 256 | :param y: 257 | a tensor of at least 2D. 258 | The last 2 dimensions will be regularized. 259 | :return: 260 | the total variation loss. 261 | """ 262 | 263 | return T.sum(T.abs(y[..., :-1] - y[..., 1:])) + T.sum(T.abs(y[..., :-1, :] - y[..., 1:, :])) 264 | 265 | 266 | def spectral_norm(module, name='weight', n_power_iterations=1, eps=1e-12, dim=None): 267 | """ 268 | Applies :func:`torch.nn.utils.spectral_norm` recursively to `module` and all of 269 | its submodules. 270 | 271 | :param module: 272 | containing module. 273 | :param name: 274 | name of weight parameter. 275 | Default: ``'weight'``. 276 | :param n_power_iterations: 277 | number of power iterations to calculate spectral norm. 278 | :param eps: 279 | epsilon for numerical stability in calculating norms. 280 | :param dim: 281 | dimension corresponding to number of outputs, 282 | the default is ``0``, except for modules that are instances of 283 | ConvTranspose{1,2,3}d, when it is ``1``. 284 | :return: 285 | the original module with the spectral norm hook. 286 | """ 287 | 288 | if hasattr(module, 'weight'): 289 | if dim is None: 290 | dim = 1 if isinstance(module, layers.ConvTranspose2d) else 0 291 | 292 | if not isinstance(module, (nn.modules.batchnorm._BatchNorm, 293 | nn.GroupNorm, 294 | nn.LayerNorm)): 295 | module = nn.utils.spectral_norm(module, name, n_power_iterations, eps, dim) 296 | 297 | return module 298 | else: 299 | for mod_name, mod in module.named_children(): 300 | mod = spectral_norm(mod, name, n_power_iterations, eps, dim) 301 | module.__setattr__(mod_name, mod) 302 | return module 303 | -------------------------------------------------------------------------------- /neuralnet_pytorch/zoo/resnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torchvision.models.resnet import model_urls 3 | from torchvision.models.utils import load_state_dict_from_url 4 | from ..layers import Sequential, Conv2d, BatchNorm2d, FC, MaxPool2d, GlobalAvgPool2D, GroupNorm, \ 5 | ResNetBottleneckBlock, ResNetBasicBlock 6 | from ..utils import batch_set_tensor 7 | 8 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 9 | 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d', 10 | 'wide_resnet50_2', 'wide_resnet101_2'] 11 | 12 | 13 | def conv1x1(in_planes, out_planes, stride=1): 14 | """1x1 convolution""" 15 | return Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) 16 | 17 | 18 | class ResNet(Sequential): 19 | 20 | def __init__(self, block, layers, num_classes=1000, zero_init_residual=False, 21 | groups=1, width_per_group=64, replace_stride_with_dilation=None, 22 | norm_layer=None, default_init=True): 23 | super(ResNet, self).__init__(input_shape=3) 24 | if norm_layer is None: 25 | norm_layer = BatchNorm2d 26 | self._norm_layer = norm_layer 27 | 28 | self.inplanes = 64 29 | self.dilation = 1 30 | if replace_stride_with_dilation is None: 31 | # each element in the tuple indicates if we should replace 32 | # the 2x2 stride with a dilated convolution instead 33 | replace_stride_with_dilation = [False, False, False] 34 | if len(replace_stride_with_dilation) != 3: 35 | raise ValueError("replace_stride_with_dilation should be None " 36 | "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) 37 | self.groups = groups 38 | self.base_width = width_per_group 39 | 40 | self.stem = Sequential(*( 41 | Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False), 42 | norm_layer(self.inplanes, activation='relu'), 43 | MaxPool2d(kernel_size=3, stride=2, padding=1) 44 | )) 45 | self.layer1 = self._make_layer(block, 64, layers[0]) 46 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2, 47 | dilate=replace_stride_with_dilation[0]) 48 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2, 49 | dilate=replace_stride_with_dilation[1]) 50 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2, 51 | dilate=replace_stride_with_dilation[2]) 52 | self.avgpool = GlobalAvgPool2D() 53 | self.fc = FC(512 * block.expansion, num_classes) 54 | 55 | if default_init: 56 | for m in self.modules(): 57 | if isinstance(m, Conv2d): 58 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 59 | elif isinstance(m, (BatchNorm2d, GroupNorm)): 60 | nn.init.constant_(m.weight, 1) 61 | nn.init.constant_(m.bias, 0) 62 | 63 | # Zero-initialize the last BN in each residual branch, 64 | # so that the residual branch starts with zeros, and each residual block behaves like an identity. 65 | # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 66 | if zero_init_residual: 67 | for m in self.modules(): 68 | if isinstance(m, ResNetBottleneckBlock): 69 | nn.init.constant_(m.bn3.weight, 0) 70 | elif isinstance(m, ResNetBasicBlock): 71 | nn.init.constant_(m.bn2.weight, 0) 72 | 73 | def _make_layer(self, block, planes, blocks, stride=1, dilate=False): 74 | norm_layer = self._norm_layer 75 | downsample = None 76 | previous_dilation = self.dilation 77 | if dilate: 78 | self.dilation *= stride 79 | stride = 1 80 | if stride != 1 or self.inplanes != planes * block.expansion: 81 | downsample = Sequential( 82 | conv1x1(self.inplanes, planes * block.expansion, stride), 83 | norm_layer(planes * block.expansion), 84 | ) 85 | 86 | layers = [] 87 | layers.append(block(self.inplanes, planes, stride=stride, downsample=downsample, groups=self.groups, 88 | base_width=self.base_width, dilation=previous_dilation, norm_layer=norm_layer)) 89 | self.inplanes = planes * block.expansion 90 | for _ in range(1, blocks): 91 | layers.append(block(self.inplanes, planes, groups=self.groups, 92 | base_width=self.base_width, dilation=self.dilation, 93 | norm_layer=norm_layer)) 94 | 95 | return Sequential(*layers) 96 | 97 | 98 | def _resnet(arch, block, layers, pretrained, progress, **kwargs): 99 | model = ResNet(block, layers, **kwargs) 100 | if pretrained: 101 | state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) 102 | try: 103 | batch_set_tensor(model.state_dict().values(), state_dict.values()) 104 | except (RuntimeError, ValueError): 105 | state_dict_iter = iter(state_dict.items()) 106 | for k, v in model.state_dict().items(): 107 | if 'num_batches_tracked' not in k: 108 | k_t, v_t = next(state_dict_iter) 109 | param_name = k.split('.')[-1] 110 | value_name = k_t.split('.') 111 | if param_name != value_name[-1]: 112 | value_name[-1] = param_name 113 | value_name = '.'.join(value_name) 114 | v.data.copy_(state_dict[value_name].data) 115 | else: 116 | v.data.copy_(v_t.data) 117 | return model 118 | 119 | 120 | def resnet18(pretrained=False, progress=True, **kwargs): 121 | """ 122 | ResNet-18 model from 123 | `"Deep Residual Learning for Image Recognition" `_ 124 | 125 | :param pretrained: 126 | If True, returns a model pre-trained on ImageNet. 127 | :param progress: 128 | If True, displays a progress bar of the download to stderr. 129 | """ 130 | return _resnet('resnet18', ResNetBasicBlock, [2, 2, 2, 2], pretrained, progress, 131 | **kwargs) 132 | 133 | 134 | def resnet34(pretrained=False, progress=True, **kwargs): 135 | """ 136 | ResNet-34 model from 137 | `"Deep Residual Learning for Image Recognition" `_ 138 | 139 | :param pretrained: 140 | If True, returns a model pre-trained on ImageNet. 141 | :param progress: 142 | If True, displays a progress bar of the download to stderr. 143 | """ 144 | return _resnet('resnet34', ResNetBasicBlock, [3, 4, 6, 3], pretrained, progress, 145 | **kwargs) 146 | 147 | 148 | def resnet50(pretrained=False, progress=True, **kwargs): 149 | r"""ResNet-50 model from 150 | `"Deep Residual Learning for Image Recognition" `_ 151 | Args: 152 | pretrained (bool): If True, returns a model pre-trained on ImageNet 153 | progress (bool): If True, displays a progress bar of the download to stderr 154 | """ 155 | return _resnet('resnet50', ResNetBottleneckBlock, [3, 4, 6, 3], pretrained, progress, 156 | **kwargs) 157 | 158 | 159 | def resnet101(pretrained=False, progress=True, **kwargs): 160 | """ 161 | ResNet-101 model from 162 | `"Deep Residual Learning for Image Recognition" `_ 163 | 164 | :param pretrained: 165 | If True, returns a model pre-trained on ImageNet. 166 | :param progress: 167 | If True, displays a progress bar of the download to stderr. 168 | """ 169 | return _resnet('resnet101', ResNetBottleneckBlock, [3, 4, 23, 3], pretrained, progress, 170 | **kwargs) 171 | 172 | 173 | def resnet152(pretrained=False, progress=True, **kwargs): 174 | """ 175 | ResNet-152 model from 176 | `"Deep Residual Learning for Image Recognition" `_ 177 | 178 | :param pretrained: 179 | If True, returns a model pre-trained on ImageNet. 180 | :param progress: 181 | If True, displays a progress bar of the download to stderr. 182 | """ 183 | return _resnet('resnet152', ResNetBottleneckBlock, [3, 8, 36, 3], pretrained, progress, 184 | **kwargs) 185 | 186 | 187 | def resnext50_32x4d(pretrained=False, progress=True, **kwargs): 188 | """ 189 | ResNeXt-50 32x4d model from 190 | `"Aggregated Residual Transformation for Deep Neural Networks" `_ 191 | 192 | :param pretrained: 193 | If True, returns a model pre-trained on ImageNet. 194 | :param progress: 195 | If True, displays a progress bar of the download to stderr. 196 | """ 197 | kwargs['groups'] = 32 198 | kwargs['width_per_group'] = 4 199 | return _resnet('resnext50_32x4d', ResNetBottleneckBlock, [3, 4, 6, 3], 200 | pretrained, progress, **kwargs) 201 | 202 | 203 | def resnext101_32x8d(pretrained=False, progress=True, **kwargs): 204 | """ 205 | ResNeXt-101 32x8d model from 206 | `"Aggregated Residual Transformation for Deep Neural Networks" `_ 207 | 208 | :param pretrained: 209 | If True, returns a model pre-trained on ImageNet. 210 | :param progress: 211 | If True, displays a progress bar of the download to stderr. 212 | """ 213 | kwargs['groups'] = 32 214 | kwargs['width_per_group'] = 8 215 | return _resnet('resnext101_32x8d', ResNetBottleneckBlock, [3, 4, 23, 3], 216 | pretrained, progress, **kwargs) 217 | 218 | 219 | def wide_resnet50_2(pretrained=False, progress=True, **kwargs): 220 | """ 221 | Wide ResNet-50-2 model from 222 | `"Wide Residual Networks" `_ 223 | The model is the same as ResNet except for the bottleneck number of channels 224 | which is twice larger in every block. The number of channels in outer 1x1 225 | convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 226 | channels, and in Wide ResNet-50-2 has 2048-1024-2048. 227 | 228 | :param pretrained: 229 | If True, returns a model pre-trained on ImageNet. 230 | :param progress: 231 | If True, displays a progress bar of the download to stderr. 232 | """ 233 | kwargs['width_per_group'] = 64 * 2 234 | return _resnet('wide_resnet50_2', ResNetBottleneckBlock, [3, 4, 6, 3], 235 | pretrained, progress, **kwargs) 236 | 237 | 238 | def wide_resnet101_2(pretrained=False, progress=True, **kwargs): 239 | """ 240 | Wide ResNet-101-2 model from 241 | `"Wide Residual Networks" `_ 242 | The model is the same as ResNet except for the bottleneck number of channels 243 | which is twice larger in every block. The number of channels in outer 1x1 244 | convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 245 | channels, and in Wide ResNet-50-2 has 2048-1024-2048. 246 | 247 | :param pretrained: 248 | If True, returns a model pre-trained on ImageNet. 249 | :param progress: 250 | If True, displays a progress bar of the download to stderr. 251 | """ 252 | kwargs['width_per_group'] = 64 * 2 253 | return _resnet('wide_resnet101_2', ResNetBottleneckBlock, [3, 4, 23, 3], 254 | pretrained, progress, **kwargs) 255 | -------------------------------------------------------------------------------- /examples/cifar.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | import torch as T 4 | import torchvision 5 | from torchvision import transforms 6 | import torch.multiprocessing as mp 7 | from torch.multiprocessing import Queue, Process 8 | import neuralnet_pytorch as nnt 9 | import neuralnet_pytorch.gin_nnt as gin 10 | from inspect import signature 11 | from glob import glob 12 | import os 13 | from packaging import version 14 | 15 | if version.parse(T.__version__) < version.parse('1.4.0'): 16 | raise ValueError('Expect Pytorch version of at least `1.4.0`') 17 | 18 | gin.enter_interactive_mode() 19 | parser = argparse.ArgumentParser(description='Neuralnet-pytorch CIFAR10 Training') 20 | parser.add_argument('config', type=str, help='Gin config file to dictate training') 21 | parser.add_argument('--test', action='store_true', help='whether to enter test mode') 22 | parser.add_argument('--device', '-d', type=str, default=['0'], nargs='+', help='device(s) to run training') 23 | parser.add_argument('--no-wait-eval', action='store_true', help='whether to use multiprocessing for evaluation') 24 | parser.add_argument('--eval-device', type=int, default=0, help='device to run evaluation') 25 | args = parser.parse_args() 26 | 27 | device = [T.device(int(dev) if dev.isdigit() else dev) for dev in args.device] 28 | eval_device = T.device(args.eval_device) 29 | no_wait_eval = args.no_wait_eval 30 | if no_wait_eval: 31 | assert eval_device not in device, 'device for evaluation must be different from the one for training' 32 | 33 | lock = nnt.utils.ReadWriteLock() 34 | if nnt.cuda_available: 35 | T.backends.cudnn.benchmark = True 36 | 37 | config_file = args.config 38 | backup_files = (__file__, config_file) 39 | 40 | transform_train = transforms.Compose([ 41 | transforms.Resize(224), 42 | transforms.RandomCrop(224, padding=32), 43 | transforms.RandomHorizontalFlip(), 44 | transforms.ToTensor(), 45 | transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), 46 | ]) 47 | transform_test = transforms.Compose([ 48 | transforms.Resize(224), 49 | transforms.ToTensor(), 50 | transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), 51 | ]) 52 | 53 | 54 | def get_loss(net, images, labels, reduction='mean'): 55 | pred = net(images) 56 | _, predicted = pred.max(1) 57 | accuracy = predicted.eq(labels).float().mean() 58 | 59 | loss = T.nn.functional.cross_entropy(pred, labels, reduction=reduction) 60 | return loss, accuracy 61 | 62 | 63 | @gin.configurable('Classifier') 64 | def train_eval(name, model, dataset, optimizer, scheduler, lr=1e-1, weight_decay=5e-4, bs=128, n_epochs=300, 65 | start_epoch=None, print_freq=1000, val_freq=10000, checkpoint_folder=None, version=-1, 66 | use_jit=True, use_amp=False, opt_level='O1', **kwargs): 67 | assert dataset in ('cifar10', 'cifar100') 68 | if use_amp: 69 | import apex 70 | 71 | net = model(num_classes=10 if dataset == 'cifar10' else 100, default_init=False) 72 | net = net.to(device[0]) 73 | 74 | opt_sig = signature(optimizer) 75 | opt_kwargs = dict([(k, kwargs[k]) for k in kwargs.keys() if k in opt_sig.parameters.keys()]) 76 | optimizer = optimizer(net.trainable, lr=lr, weight_decay=weight_decay, **opt_kwargs) 77 | if scheduler is not None: 78 | sch_sig = signature(scheduler) 79 | sch_kwargs = dict([(k, kwargs[k]) for k in kwargs.keys() if k in sch_sig.parameters.keys()]) 80 | scheduler = scheduler(optimizer, **sch_kwargs) 81 | 82 | dataset_ = torchvision.datasets.CIFAR10 if dataset == 'cifar10' else torchvision.datasets.CIFAR100 83 | train_data = dataset_(root='./data', train=True, download=True, transform=transform_train) 84 | train_loader = T.utils.data.DataLoader(train_data, batch_size=bs, shuffle=True, num_workers=5) 85 | 86 | if checkpoint_folder is None: 87 | mon = nnt.Monitor(name, print_freq=print_freq, num_iters=int(np.ceil(len(train_data) / bs)), 88 | use_tensorboard=True) 89 | mon.backup(backup_files) 90 | 91 | mon.dump_rep('network', net) 92 | mon.dump_rep('optimizer', optimizer) 93 | 94 | states = { 95 | 'model_state_dict': net.state_dict(), 96 | 'opt_state_dict': optimizer.state_dict() 97 | } 98 | 99 | if scheduler is not None: 100 | mon.dump_rep('scheduler', scheduler) 101 | states['scheduler_state_dict'] = scheduler.state_dict() 102 | 103 | else: 104 | mon = nnt.Monitor(current_folder=checkpoint_folder, print_freq=print_freq, num_iters=len(train_data) // bs, 105 | use_tensorboard=True) 106 | states = mon.load('training.pt', method='torch', version=version) 107 | net.load_state_dict(states['model_state_dict']) 108 | optimizer.load_state_dict(states['opt_state_dict']) 109 | if scheduler: 110 | scheduler.load_state_dict(states['scheduler_state_dict']) 111 | 112 | if use_amp and 'amp' in states.keys(): 113 | apex.amp.load_state_dict(states['amp']) 114 | 115 | if start_epoch: 116 | start_epoch = start_epoch - 1 117 | mon.epoch = start_epoch 118 | 119 | print('Resume from epoch %d...' % mon.epoch) 120 | 121 | if not no_wait_eval: 122 | eval_data = dataset_(root='./data', train=False, download=True, transform=transform_test) 123 | eval_loader = T.utils.data.DataLoader(eval_data, batch_size=bs, shuffle=False, num_workers=2) 124 | 125 | if nnt.cuda_available: 126 | train_loader = nnt.DataPrefetcher(train_loader, device=device[0]) 127 | if not no_wait_eval: 128 | eval_loader = nnt.DataPrefetcher(eval_loader, device=device[0]) 129 | 130 | if use_jit: 131 | img = T.rand(1, 3, 32, 32).to(device[0]) 132 | net.train(True) 133 | net_train = T.jit.trace(net, img) 134 | net.eval() 135 | net_eval = T.jit.trace(net, img) 136 | 137 | if use_amp: 138 | if use_jit: 139 | net_train, optimizer = apex.amp.initialize(net_train, optimizer, opt_level=opt_level) 140 | net_eval = apex.amp.initialize(net_eval, opt_level=opt_level) 141 | else: 142 | net, optimizer = apex.amp.initialize(net, optimizer, opt_level=opt_level) 143 | 144 | if 'amp' not in states.keys(): 145 | states['amp'] = apex.amp.state_dict() 146 | 147 | if use_jit: 148 | net_train = T.nn.DataParallel(net_train, device_ids=device) 149 | net_eval = T.nn.DataParallel(net_eval, device_ids=device) 150 | else: 151 | net = T.nn.DataParallel(net, device_ids=device) 152 | 153 | def learn(images, labels, reduction='mean'): 154 | net.train(True) 155 | optimizer.zero_grad() 156 | loss, accuracy = get_loss(net_train if use_jit else net, images, labels, reduction=reduction) 157 | if not (T.isnan(loss) or T.isinf(loss)): 158 | if use_amp: 159 | with apex.amp.scale_loss(loss, optimizer) as scaled_loss: 160 | scaled_loss.backward() 161 | else: 162 | loss.backward() 163 | optimizer.step() 164 | else: 165 | raise ValueError('NaN encountered!') 166 | 167 | mon.plot('train-loss', nnt.utils.to_numpy(loss), smooth=.99) 168 | mon.plot('train-accuracy', nnt.utils.to_numpy(accuracy), smooth=.99) 169 | del loss, accuracy 170 | 171 | if no_wait_eval: 172 | q = Queue() 173 | eval_proc = Process(target=eval_queue, 174 | args=(q, mon.current_folder, dataset, bs, use_jit, use_amp, opt_level)) 175 | eval_proc.daemon = True 176 | eval_proc.start() 177 | 178 | start_epoch = mon.epoch if start_epoch is None else start_epoch 179 | print('Training...') 180 | with T.jit.optimized_execution(use_jit): 181 | for _ in mon.iter_epoch(range(start_epoch, n_epochs)): 182 | for idx, lr_ in enumerate(scheduler.get_last_lr()): 183 | mon.plot('lr-%d' % idx, lr_, filter_outliers=False) 184 | 185 | for batch in mon.iter_batch(train_loader): 186 | batch = nnt.utils.batch_to_device(batch, device[0]) 187 | 188 | learn(*batch) 189 | if val_freq and mon.iter % val_freq == 0: 190 | if no_wait_eval: 191 | lock.acquire_write() 192 | mon.dump('tmp.pt', states, method='torch') 193 | lock.release_write() 194 | q.put((mon.epoch, mon.iter)) 195 | q.put(None) 196 | else: 197 | net.eval() 198 | with T.set_grad_enabled(False): 199 | losses, accuracies = [], [] 200 | for itt, batch in enumerate(eval_loader): 201 | batch = nnt.utils.batch_to_device(batch, device[0]) 202 | 203 | loss, acc = get_loss(net_eval if use_jit else net, *batch) 204 | losses.append(nnt.utils.to_numpy(loss)) 205 | accuracies.append(nnt.utils.to_numpy(acc)) 206 | 207 | mon.plot('test-loss', np.mean(losses)) 208 | mon.plot('test-accuracy', np.mean(accuracies)) 209 | mon.dump('training.pt', states, method='torch', keep=10) 210 | if scheduler is not None: 211 | scheduler.step() 212 | 213 | if no_wait_eval: 214 | q.put('DONE') 215 | eval_proc.join() 216 | 217 | print('Training finished!') 218 | 219 | 220 | def eval_queue(q, ckpt, dataset, bs, use_jit, use_amp, opt_level): 221 | assert dataset in ('cifar10', 'cifar100') 222 | 223 | # TODO: passing `model` this way is ugly 224 | @gin.configurable('Classifier') 225 | def _make_network(model, **kwargs): 226 | net = model(num_classes=10 if dataset == 'cifar10' else 100, default_init=False) 227 | net = net.to(eval_device) 228 | net.eval() 229 | return net 230 | 231 | lock.acquire_read() 232 | mon = nnt.Monitor(current_folder=ckpt) 233 | lock.release_read() 234 | 235 | cfg = glob(os.path.join(mon.file_folder, '*.gin')) 236 | cfg = cfg[0] if len(cfg) == 1 else config_file # fall back 237 | gin.parse_config_file(cfg) 238 | net = _make_network() 239 | 240 | dataset = torchvision.datasets.CIFAR10 if dataset == 'cifar10' else torchvision.datasets.CIFAR100 241 | eval_data = dataset(root='./data', train=False, download=True, transform=transform_test) 242 | eval_loader = T.utils.data.DataLoader(eval_data, batch_size=bs, shuffle=False) 243 | 244 | if nnt.cuda_available: 245 | eval_loader = nnt.DataPrefetcher(eval_loader, device=eval_device) 246 | 247 | while True: 248 | item = q.get() 249 | if item == 'DONE': 250 | break 251 | elif item is not None: 252 | lock.acquire_read() 253 | mon.load_state() 254 | lock.release_read() 255 | mon.epoch, mon.iter = item 256 | # TODO: find a better way to load the current state of `monitor` 257 | lock.acquire_read() 258 | states = mon.load('tmp.pt', method='torch') 259 | lock.release_read() 260 | net.load_state_dict(states['model_state_dict']) 261 | if use_jit: 262 | img = T.rand(1, 3, 32, 32).to(eval_device) 263 | net = T.jit.trace(net, img) 264 | 265 | if use_amp: 266 | import apex 267 | net = apex.amp.initialize(net, opt_level=opt_level) 268 | apex.amp.load_state_dict(states['amp']) 269 | 270 | with T.set_grad_enabled(False): 271 | losses, accuracies = [], [] 272 | for itt, batch in enumerate(eval_loader): 273 | batch = nnt.utils.batch_to_device(batch, eval_device) 274 | 275 | loss, acc = get_loss(net, *batch) 276 | losses.append(nnt.utils.to_numpy(loss)) 277 | accuracies.append(nnt.utils.to_numpy(acc)) 278 | 279 | mon.plot('test-loss', np.mean(losses)) 280 | mon.plot('test-accuracy', np.mean(accuracies)) 281 | mon.flush() 282 | 283 | 284 | @gin.configurable('Classifier') 285 | def test(name, model, dataset, bs=128, print_freq=1000, checkpoint_folder=None, version=-1, 286 | use_jit=True, use_amp=False, opt_level='O1', **kwargs): 287 | assert dataset in ('cifar10', 'cifar100') 288 | 289 | net = model(num_classes=10 if dataset == 'cifar10' else 100, default_init=False) 290 | net = net.to(device[0]) 291 | net.eval() 292 | 293 | dataset = torchvision.datasets.CIFAR10 if dataset == 'cifar10' else torchvision.datasets.CIFAR100 294 | test_data = dataset(root='./data', train=False, download=True, transform=transform_test) 295 | test_loader = T.utils.data.DataLoader(test_data, batch_size=bs, shuffle=False, num_workers=5) 296 | 297 | mon = nnt.Monitor(current_folder=checkpoint_folder, print_freq=print_freq, num_iters=len(test_data) // bs, 298 | use_tensorboard=True) 299 | states = mon.load('training.pt', method='torch', version=version) 300 | net.load_state_dict(states['model_state_dict']) 301 | if use_jit: 302 | img = T.rand(1, 3, 32, 32).to(device[0]) 303 | net = T.jit.trace(net, img) 304 | 305 | if use_amp: 306 | import apex 307 | net = apex.amp.initialize(net, opt_level=opt_level) 308 | apex.amp.load_state_dict(states['amp']) 309 | 310 | net = T.nn.DataParallel(net, device_ids=device) 311 | print('Resume from epoch %d...' % mon.epoch) 312 | print('Testing...') 313 | with T.set_grad_enabled(False): 314 | for itt, batch in mon.iter_batch(enumerate(test_loader)): 315 | batch = nnt.utils.batch_to_device(batch, device[0]) 316 | 317 | loss, acc = get_loss(net,*batch) 318 | mon.plot('test-loss', nnt.utils.to_numpy(loss)) 319 | mon.plot('test-accuracy', nnt.utils.to_numpy(acc)) 320 | print('Testing finished!') 321 | 322 | 323 | if __name__ == '__main__': 324 | mp.set_start_method('spawn') 325 | gin.parse_config_file(config_file) 326 | test() if args.test else train_eval() 327 | -------------------------------------------------------------------------------- /neuralnet_pytorch/extensions/csrc/emd_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include // at::cuda::getApplyGrid 3 | #include 4 | #include 5 | #include 6 | 7 | /******************************** 8 | * Forward kernel for approxmatch 9 | *********************************/ 10 | 11 | template 12 | __global__ void 13 | approxmatch(int b, int n, int m, const scalar_t* __restrict__ xyz1, 14 | const scalar_t* __restrict__ xyz2, scalar_t* __restrict__ match, 15 | scalar_t* temp) 16 | { 17 | scalar_t *remainL = temp + blockIdx.x * (n + m) * 2, 18 | *remainR = temp + blockIdx.x * (n + m) * 2 + n, 19 | *ratioL = temp + blockIdx.x * (n + m) * 2 + n + m, 20 | *ratioR = temp + blockIdx.x * (n + m) * 2 + n + m + n; 21 | scalar_t multiL, multiR; 22 | if (n >= m) { 23 | multiL = 1; 24 | multiR = n / m; 25 | } else { 26 | multiL = m / n; 27 | multiR = 1; 28 | } 29 | const int Block = 1024; 30 | __shared__ scalar_t buf[Block * 4]; 31 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 32 | for (int j = threadIdx.x; j < n * m; j += blockDim.x) 33 | match[i * n * m + j] = 0; 34 | for (int j = threadIdx.x; j < n; j += blockDim.x) 35 | remainL[j] = multiL; 36 | for (int j = threadIdx.x; j < m; j += blockDim.x) 37 | remainR[j] = multiR; 38 | __syncthreads(); 39 | for (int j = 7; j >= -2; j--) { 40 | scalar_t level = -powf(4.0f, j); 41 | if (j == -2) { 42 | level = 0; 43 | } 44 | for (int k0 = 0; k0 < n; k0 += blockDim.x) { 45 | int k = k0 + threadIdx.x; 46 | scalar_t x1 = 0, y1 = 0, z1 = 0; 47 | if (k < n) { 48 | x1 = xyz1[i * n * 3 + k * 3 + 0]; 49 | y1 = xyz1[i * n * 3 + k * 3 + 1]; 50 | z1 = xyz1[i * n * 3 + k * 3 + 2]; 51 | } 52 | scalar_t suml = 1e-9f; 53 | for (int l0 = 0; l0 < m; l0 += Block) { 54 | int lend = min(m, l0 + Block) - l0; 55 | for (int l = threadIdx.x; l < lend; l += blockDim.x) { 56 | scalar_t x2 = xyz2[i * m * 3 + l0 * 3 + l * 3 + 0]; 57 | scalar_t y2 = xyz2[i * m * 3 + l0 * 3 + l * 3 + 1]; 58 | scalar_t z2 = xyz2[i * m * 3 + l0 * 3 + l * 3 + 2]; 59 | buf[l * 4 + 0] = x2; 60 | buf[l * 4 + 1] = y2; 61 | buf[l * 4 + 2] = z2; 62 | buf[l * 4 + 3] = remainR[l0 + l]; 63 | } 64 | __syncthreads(); 65 | for (int l = 0; l < lend; l++) { 66 | scalar_t x2 = buf[l * 4 + 0]; 67 | scalar_t y2 = buf[l * 4 + 1]; 68 | scalar_t z2 = buf[l * 4 + 2]; 69 | scalar_t d = 70 | level * ((x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + 71 | (z2 - z1) * (z2 - z1)); 72 | scalar_t w = __expf(d) * buf[l * 4 + 3]; 73 | suml += w; 74 | } 75 | __syncthreads(); 76 | } 77 | if (k < n) 78 | ratioL[k] = remainL[k] / suml; 79 | } 80 | __syncthreads(); 81 | for (int l0 = 0; l0 < m; l0 += blockDim.x) { 82 | int l = l0 + threadIdx.x; 83 | scalar_t x2 = 0, y2 = 0, z2 = 0; 84 | if (l < m) { 85 | x2 = xyz2[i * m * 3 + l * 3 + 0]; 86 | y2 = xyz2[i * m * 3 + l * 3 + 1]; 87 | z2 = xyz2[i * m * 3 + l * 3 + 2]; 88 | } 89 | scalar_t sumr = 0; 90 | for (int k0 = 0; k0 < n; k0 += Block) { 91 | int kend = min(n, k0 + Block) - k0; 92 | for (int k = threadIdx.x; k < kend; k += blockDim.x) { 93 | buf[k * 4 + 0] = xyz1[i * n * 3 + k0 * 3 + k * 3 + 0]; 94 | buf[k * 4 + 1] = xyz1[i * n * 3 + k0 * 3 + k * 3 + 1]; 95 | buf[k * 4 + 2] = xyz1[i * n * 3 + k0 * 3 + k * 3 + 2]; 96 | buf[k * 4 + 3] = ratioL[k0 + k]; 97 | } 98 | __syncthreads(); 99 | for (int k = 0; k < kend; k++) { 100 | scalar_t x1 = buf[k * 4 + 0]; 101 | scalar_t y1 = buf[k * 4 + 1]; 102 | scalar_t z1 = buf[k * 4 + 2]; 103 | scalar_t w = 104 | __expf(level * ((x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + 105 | (z2 - z1) * (z2 - z1))) * 106 | buf[k * 4 + 3]; 107 | sumr += w; 108 | } 109 | __syncthreads(); 110 | } 111 | if (l < m) { 112 | sumr *= remainR[l]; 113 | scalar_t consumption = fminf(remainR[l] / (sumr + 1e-9f), 1.0f); 114 | ratioR[l] = consumption * remainR[l]; 115 | remainR[l] = fmaxf(scalar_t(0.0), remainR[l] - sumr); 116 | } 117 | } 118 | __syncthreads(); 119 | for (int k0 = 0; k0 < n; k0 += blockDim.x) { 120 | int k = k0 + threadIdx.x; 121 | scalar_t x1 = 0, y1 = 0, z1 = 0; 122 | if (k < n) { 123 | x1 = xyz1[i * n * 3 + k * 3 + 0]; 124 | y1 = xyz1[i * n * 3 + k * 3 + 1]; 125 | z1 = xyz1[i * n * 3 + k * 3 + 2]; 126 | } 127 | scalar_t suml = 0; 128 | for (int l0 = 0; l0 < m; l0 += Block) { 129 | int lend = min(m, l0 + Block) - l0; 130 | for (int l = threadIdx.x; l < lend; l += blockDim.x) { 131 | buf[l * 4 + 0] = xyz2[i * m * 3 + l0 * 3 + l * 3 + 0]; 132 | buf[l * 4 + 1] = xyz2[i * m * 3 + l0 * 3 + l * 3 + 1]; 133 | buf[l * 4 + 2] = xyz2[i * m * 3 + l0 * 3 + l * 3 + 2]; 134 | buf[l * 4 + 3] = ratioR[l0 + l]; 135 | } 136 | __syncthreads(); 137 | scalar_t rl = ratioL[k]; 138 | if (k < n) { 139 | for (int l = 0; l < lend; l++) { 140 | scalar_t x2 = buf[l * 4 + 0]; 141 | scalar_t y2 = buf[l * 4 + 1]; 142 | scalar_t z2 = buf[l * 4 + 2]; 143 | scalar_t w = 144 | __expf(level * ((x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + 145 | (z2 - z1) * (z2 - z1))) * 146 | rl * buf[l * 4 + 3]; 147 | match[i * n * m + (l0 + l) * n + k] += w; 148 | suml += w; 149 | } 150 | } 151 | __syncthreads(); 152 | } 153 | if (k < n) 154 | remainL[k] = fmaxf(scalar_t(0.0), remainL[k] - suml); 155 | } 156 | __syncthreads(); 157 | } 158 | } 159 | } 160 | 161 | /* ApproxMatch forward interface 162 | Input: 163 | xyz1: (B, N1, 3) # dataset_points 164 | xyz2: (B, N2, 3) # query_points 165 | Output: 166 | match: (B, N2, N1) 167 | */ 168 | at::Tensor 169 | approx_match_cuda_forward(const at::Tensor xyz1, const at::Tensor xyz2) 170 | { 171 | const auto b = xyz1.size(0); 172 | const auto n = xyz1.size(1); 173 | const auto m = xyz2.size(1); 174 | 175 | auto match = at::zeros({ b, m, n }, xyz1.type()); 176 | auto temp = at::zeros({ b, (n + m) * 2 }, xyz1.type()); 177 | 178 | AT_DISPATCH_FLOATING_TYPES_AND( 179 | at::ScalarType::Half, xyz1.scalar_type(), "approx_match_forward_cuda", 180 | ([&] { 181 | approxmatch<<<32, 512>>>( 182 | b, n, m, xyz1.data(), xyz2.data(), 183 | match.data(), temp.data()); 184 | })); 185 | THCudaCheck(cudaGetLastError()); 186 | 187 | return match; 188 | } 189 | 190 | /******************************** 191 | * Forward kernel for matchcost 192 | *********************************/ 193 | 194 | template 195 | __global__ void 196 | matchcost(int b, int n, int m, const scalar_t* __restrict__ xyz1, 197 | const scalar_t* __restrict__ xyz2, const scalar_t* __restrict__ match, 198 | scalar_t* __restrict__ out) 199 | { 200 | __shared__ scalar_t allsum[512]; 201 | const int Block = 1024; 202 | __shared__ scalar_t buf[Block * 3]; 203 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 204 | scalar_t subsum = 0; 205 | for (int k0 = 0; k0 < n; k0 += blockDim.x) { 206 | int k = k0 + threadIdx.x; 207 | scalar_t x1 = 0, y1 = 0, z1 = 0; 208 | if (k < n) { 209 | x1 = xyz1[i * n * 3 + k * 3 + 0]; 210 | y1 = xyz1[i * n * 3 + k * 3 + 1]; 211 | z1 = xyz1[i * n * 3 + k * 3 + 2]; 212 | } 213 | for (int l0 = 0; l0 < m; l0 += Block) { 214 | int lend = min(m, l0 + Block) - l0; 215 | for (int l = threadIdx.x; l < lend * 3; l += blockDim.x) 216 | buf[l] = xyz2[i * m * 3 + l0 * 3 + l]; 217 | __syncthreads(); 218 | if (k < n) { 219 | for (int l = 0; l < lend; l++) { 220 | scalar_t x2 = buf[l * 3 + 0]; 221 | scalar_t y2 = buf[l * 3 + 1]; 222 | scalar_t z2 = buf[l * 3 + 2]; 223 | scalar_t d = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + 224 | (z2 - z1) * (z2 - z1); 225 | subsum += d * match[i * n * m + (l0 + l) * n + k]; 226 | } 227 | } 228 | __syncthreads(); 229 | } 230 | } 231 | allsum[threadIdx.x] = subsum; 232 | for (int j = 1; j < blockDim.x; j <<= 1) { 233 | __syncthreads(); 234 | if ((threadIdx.x & j) == 0 && threadIdx.x + j < blockDim.x) { 235 | allsum[threadIdx.x] += allsum[threadIdx.x + j]; 236 | } 237 | } 238 | if (threadIdx.x == 0) 239 | out[i] = allsum[0]; 240 | __syncthreads(); 241 | } 242 | } 243 | 244 | /* MatchCost forward interface 245 | Input: 246 | xyz1: (B, N1, 3) # dataset_points 247 | xyz2: (B, N2, 3) # query_points 248 | match: (B, N2, N1) 249 | Output: 250 | cost: (B) 251 | */ 252 | at::Tensor 253 | match_cost_cuda_forward(const at::Tensor xyz1, const at::Tensor xyz2, 254 | const at::Tensor match) 255 | { 256 | const auto b = xyz1.size(0); 257 | const auto n = xyz1.size(1); 258 | const auto m = xyz2.size(1); 259 | 260 | auto cost = at::zeros({ b }, xyz1.type()); 261 | 262 | AT_DISPATCH_FLOATING_TYPES_AND( 263 | at::ScalarType::Half, xyz1.scalar_type(), "match_cost_forward_cuda", ([&] { 264 | matchcost<<<32, 512>>>( 265 | b, n, m, xyz1.data(), xyz2.data(), 266 | match.data(), cost.data()); 267 | })); 268 | THCudaCheck(cudaGetLastError()); 269 | 270 | return cost; 271 | } 272 | 273 | /******************************** 274 | * matchcostgrad2 kernel 275 | *********************************/ 276 | 277 | template 278 | __global__ void 279 | matchcostgrad2(int b, int n, int m, const scalar_t* __restrict__ grad_cost, 280 | const scalar_t* __restrict__ xyz1, 281 | const scalar_t* __restrict__ xyz2, 282 | const scalar_t* __restrict__ match, scalar_t* __restrict__ grad2) 283 | { 284 | __shared__ scalar_t sum_grad[256 * 3]; 285 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 286 | int kbeg = m * blockIdx.y / gridDim.y; 287 | int kend = m * (blockIdx.y + 1) / gridDim.y; 288 | for (int k = kbeg; k < kend; k++) { 289 | scalar_t x2 = xyz2[(i * m + k) * 3 + 0]; 290 | scalar_t y2 = xyz2[(i * m + k) * 3 + 1]; 291 | scalar_t z2 = xyz2[(i * m + k) * 3 + 2]; 292 | scalar_t subsumx = 0, subsumy = 0, subsumz = 0; 293 | for (int j = threadIdx.x; j < n; j += blockDim.x) { 294 | scalar_t x1 = x2 - xyz1[(i * n + j) * 3 + 0]; 295 | scalar_t y1 = y2 - xyz1[(i * n + j) * 3 + 1]; 296 | scalar_t z1 = z2 - xyz1[(i * n + j) * 3 + 2]; 297 | scalar_t d = match[i * n * m + k * n + j] * 2; 298 | subsumx += x1 * d; 299 | subsumy += y1 * d; 300 | subsumz += z1 * d; 301 | } 302 | sum_grad[threadIdx.x * 3 + 0] = subsumx; 303 | sum_grad[threadIdx.x * 3 + 1] = subsumy; 304 | sum_grad[threadIdx.x * 3 + 2] = subsumz; 305 | for (int j = 1; j < blockDim.x; j <<= 1) { 306 | __syncthreads(); 307 | int j1 = threadIdx.x; 308 | int j2 = threadIdx.x + j; 309 | if ((j1 & j) == 0 && j2 < blockDim.x) { 310 | sum_grad[j1 * 3 + 0] += sum_grad[j2 * 3 + 0]; 311 | sum_grad[j1 * 3 + 1] += sum_grad[j2 * 3 + 1]; 312 | sum_grad[j1 * 3 + 2] += sum_grad[j2 * 3 + 2]; 313 | } 314 | } 315 | if (threadIdx.x == 0) { 316 | grad2[(i * m + k) * 3 + 0] = sum_grad[0] * grad_cost[i]; 317 | grad2[(i * m + k) * 3 + 1] = sum_grad[1] * grad_cost[i]; 318 | grad2[(i * m + k) * 3 + 2] = sum_grad[2] * grad_cost[i]; 319 | } 320 | __syncthreads(); 321 | } 322 | } 323 | } 324 | 325 | /******************************** 326 | * matchcostgrad1 kernel 327 | *********************************/ 328 | 329 | template 330 | __global__ void 331 | matchcostgrad1(int b, int n, int m, const scalar_t* __restrict__ grad_cost, 332 | const scalar_t* __restrict__ xyz1, 333 | const scalar_t* __restrict__ xyz2, 334 | const scalar_t* __restrict__ match, scalar_t* __restrict__ grad1) 335 | { 336 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 337 | for (int l = threadIdx.x; l < n; l += blockDim.x) { 338 | scalar_t x1 = xyz1[i * n * 3 + l * 3 + 0]; 339 | scalar_t y1 = xyz1[i * n * 3 + l * 3 + 1]; 340 | scalar_t z1 = xyz1[i * n * 3 + l * 3 + 2]; 341 | scalar_t dx = 0, dy = 0, dz = 0; 342 | for (int k = 0; k < m; k++) { 343 | scalar_t x2 = xyz2[i * m * 3 + k * 3 + 0]; 344 | scalar_t y2 = xyz2[i * m * 3 + k * 3 + 1]; 345 | scalar_t z2 = xyz2[i * m * 3 + k * 3 + 2]; 346 | scalar_t d = match[i * n * m + k * n + l] * 2; 347 | dx += (x1 - x2) * d; 348 | dy += (y1 - y2) * d; 349 | dz += (z1 - z2) * d; 350 | } 351 | grad1[i * n * 3 + l * 3 + 0] = dx * grad_cost[i]; 352 | grad1[i * n * 3 + l * 3 + 1] = dy * grad_cost[i]; 353 | grad1[i * n * 3 + l * 3 + 2] = dz * grad_cost[i]; 354 | } 355 | } 356 | } 357 | 358 | /* MatchCost backward interface 359 | Input: 360 | grad_cost: (B) # gradients on cost 361 | xyz1: (B, N1, 3) # dataset_points 362 | xyz2: (B, N2, 3) # query_points 363 | match: (B, N2, N1) 364 | Output: 365 | grad1: (B, N1, 3) 366 | grad2: (B, N2, 3) 367 | */ 368 | std::vector 369 | match_cost_cuda_backward(const at::Tensor grad_cost, const at::Tensor xyz1, 370 | const at::Tensor xyz2, const at::Tensor match) 371 | { 372 | const auto b = xyz1.size(0); 373 | const auto n = xyz1.size(1); 374 | const auto m = xyz2.size(1); 375 | 376 | auto grad1 = at::zeros({ b, n, 3 }, xyz1.type()); 377 | auto grad2 = at::zeros({ b, m, 3 }, xyz2.type()); 378 | 379 | AT_DISPATCH_FLOATING_TYPES_AND( 380 | at::ScalarType::Half, xyz1.scalar_type(), "match_cost_backward_cuda", ([&] { 381 | matchcostgrad1<<<32, 512>>>( 382 | b, n, m, grad_cost.data(), xyz1.data(), 383 | xyz2.data(), match.data(), grad1.data()); 384 | matchcostgrad2<<>>( 385 | b, n, m, grad_cost.data(), xyz1.data(), 386 | xyz2.data(), match.data(), grad2.data()); 387 | })); 388 | THCudaCheck(cudaGetLastError()); 389 | 390 | return { grad1, grad2 }; 391 | } 392 | -------------------------------------------------------------------------------- /neuralnet_pytorch/layers/points.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | import numpy as np 3 | import torch as T 4 | import torch.nn as nn 5 | 6 | from .abstract import wrapper, Module, Sequential 7 | from .convolution import Conv2d, DepthwiseSepConv2D, FC 8 | from .. import utils 9 | from ..utils import _pointset_shape 10 | 11 | __all__ = ['XConv', 'GraphConv', 'BatchGraphConv', 'GraphXConv'] 12 | 13 | 14 | @utils.add_custom_repr 15 | class XConv(Module): 16 | """ 17 | Performs X-Convolution on unordered set as in `this paper`_. 18 | 19 | .. _this paper: https://papers.nips.cc/paper/7362-pointcnn-convolution-on-x-transformed-points.pdf 20 | 21 | Parameters 22 | ---------- 23 | input_shape 24 | shape of the input tensor. 25 | If an integer is passed, it is treated as the size of each input sample. 26 | feature_dim : int 27 | dimension of the input features. 28 | out_channels : int 29 | number of channels produced by the convolution. 30 | out_features : int 31 | size of each output sample. 32 | num_neighbors : int 33 | size of the convolving kernel. 34 | depth_mul 35 | depth multiplier for intermediate result of depthwise convolution 36 | activation 37 | non-linear function to activate the linear result. 38 | It accepts any callable function 39 | as well as a recognizable ``str``. 40 | A list of possible ``str`` is in :func:`~neuralnet_pytorch.utils.function`. 41 | dropout : bool 42 | whether to use dropout. 43 | bn : bool 44 | whether to use batch normalization. 45 | kwargs 46 | extra keyword arguments to pass to activation. 47 | """ 48 | def __init__(self, input_shape, feature_dim, out_channels, out_features, num_neighbors, depth_mul, 49 | activation='relu', dropout=None, bn=True, **kwargs): 50 | input_shape = _pointset_shape(input_shape) 51 | 52 | super().__init__(input_shape) 53 | self.feature_dim = feature_dim 54 | self.out_channels = out_channels 55 | self.num_neighbors = num_neighbors 56 | self.out_features = out_features 57 | self.depth_mul = depth_mul 58 | self.activation = utils.function(activation, **kwargs) 59 | self.dropout = dropout 60 | self.bn = bn 61 | 62 | self.fcs = Sequential(input_shape=input_shape) 63 | self.fcs.add_module('fc1', FC(self.fcs.output_shape, out_features, activation=activation)) 64 | if dropout: 65 | self.fcs.add_module('dropout1', wrapper(self.output_shape, nn.Dropout2d, p=dropout)) 66 | self.fcs.add_module('fc2', FC(self.fcs.output_shape, out_features, activation=activation)) 67 | if dropout: 68 | self.fcs.add_module('dropout2', wrapper(self.output_shape, nn.Dropout2d, p=dropout)) 69 | 70 | from neuralnet_pytorch.layers.resizing import DimShuffle 71 | from neuralnet_pytorch.layers.normalization import BatchNorm2d 72 | 73 | self.x_trans = Sequential(input_shape=input_shape[:2] + (num_neighbors, input_shape[-1])) 74 | self.x_trans.add_module('dimshuffle1', DimShuffle(self.x_trans.output_shape, (0, 3, 1, 2))) 75 | self.x_trans.add_module('conv', Conv2d(self.x_trans.output_shape, num_neighbors ** 2, (1, num_neighbors), 76 | activation=activation, padding='valid', **kwargs)) 77 | self.x_trans.add_module('dimshuffle2', DimShuffle(self.x_trans.output_shape, (0, 2, 3, 1))) 78 | self.x_trans.add_module('fc1', FC(self.x_trans.output_shape, num_neighbors ** 2, activation='relu', **kwargs)) 79 | self.x_trans.add_module('fc2', FC(self.x_trans.output_shape, num_neighbors ** 2, **kwargs)) 80 | 81 | self.end_conv = Sequential(input_shape=input_shape[:2] + (num_neighbors, feature_dim + out_features)) 82 | self.end_conv.add_module('dimshuffle1', DimShuffle(self.end_conv.output_shape, (0, 3, 1, 2))) 83 | self.end_conv.add_module('conv', 84 | DepthwiseSepConv2D(self.end_conv.output_shape, out_channels, (1, num_neighbors), 85 | depth_mul=depth_mul, activation=None if bn else activation, 86 | padding='valid', **kwargs)) 87 | if bn: 88 | self.end_conv.add_module('bn', BatchNorm2d(self.end_conv.output_shape, momentum=.9, activation=activation)) 89 | self.end_conv.add_module('dimshuffle2', DimShuffle(self.end_conv.output_shape, (0, 2, 3, 1))) 90 | 91 | def forward(self, *input, **kwargs): 92 | rep_pt, pts, fts = input 93 | 94 | if fts is not None: 95 | assert rep_pt.size()[0] == pts.size()[0] == fts.size()[0] # Check N is equal. 96 | assert rep_pt.size()[1] == pts.size()[1] == fts.size()[1] # Check P is equal. 97 | assert pts.size()[2] == fts.size()[2] == self.num_neighbors # Check K is equal. 98 | assert fts.size()[3] == self.feature_dim # Check C_in is equal. 99 | else: 100 | assert rep_pt.size()[0] == pts.size()[0] # Check N is equal. 101 | assert rep_pt.size()[1] == pts.size()[1] # Check P is equal. 102 | assert pts.size()[2] == self.num_neighbors # Check K is equal. 103 | assert rep_pt.size()[2] == pts.size()[3] == self.input_shape[-1] # Check dims is equal. 104 | 105 | N = len(pts) 106 | P = rep_pt.size()[1] # (N, P, K, dims) 107 | p_center = T.unsqueeze(rep_pt, dim=2) # (N, P, 1, dims) 108 | 109 | # Move pts to local coordinate system of rep_pt. 110 | pts_local = pts - p_center # (N, P, K, dims) 111 | 112 | # Individually lift each point into C_mid space. 113 | fts_lifted = self.fcs(pts_local) # (N, P, K, C_mid) 114 | 115 | if fts is None: 116 | fts_cat = fts_lifted 117 | else: 118 | fts_cat = T.cat((fts_lifted, fts), -1) # (N, P, K, C_mid + C_in) 119 | 120 | # Learn the (N, K, K) X-transformation matrix. 121 | X_shape = (N, P, self.num_neighbors, self.num_neighbors) 122 | X = self.x_trans(pts_local) 123 | X = X.view(*X_shape) 124 | 125 | # Weight and permute fts_cat with the learned X. 126 | fts_X = T.matmul(X, fts_cat) 127 | fts_p = self.end_conv(fts_X).squeeze(dim=2) 128 | return fts_p 129 | 130 | @property 131 | @utils.validate 132 | def output_shape(self): 133 | return self.input_shape[:2] + (self.out_channels,) 134 | 135 | def extra_repr(self): 136 | s = ('{input_shape}, feature_dim={feature_dim}, out_channels={out_channels}, out_features={out_features}, ' 137 | 'num_neighbors={num_neighbors}, depth_mul={depth_mul}, dropout={dropout}, bn={bn}') 138 | 139 | s = s.format(**self.__dict__) 140 | s += ', activation={}'.format(self.activation.__name__) 141 | return s 142 | 143 | 144 | class GraphConv(FC): 145 | """ 146 | Performs graph convolution as described in https://arxiv.org/abs/1609.02907. 147 | Adapted from https://github.com/tkipf/pygcn/blob/master/pygcn/layers.py. 148 | 149 | Parameters 150 | ---------- 151 | input_shape 152 | shape of the input tensor. 153 | If an integer is passed, it is treated as the size of each input sample. 154 | out_features : int 155 | size of each output sample. 156 | activation 157 | non-linear function to activate the linear result. 158 | It accepts any callable function 159 | as well as a recognizable ``str``. 160 | A list of possible ``str`` is in :func:`~neuralnet_pytorch.utils.function`. 161 | kwargs 162 | extra keyword arguments to pass to activation. 163 | """ 164 | 165 | def __init__(self, input_shape, out_features, bias=True, activation=None, **kwargs): 166 | super().__init__(input_shape, out_features, bias, activation=activation, **kwargs) 167 | 168 | def reset_parameters(self): 169 | if self.weights_init is None: 170 | stdv = 1. / np.sqrt(self.weight.size(1)) 171 | self.weight.data.uniform_(-stdv, stdv) 172 | else: 173 | self.weights_init(self.weight) 174 | 175 | if self.bias is not None: 176 | if self.bias_init is None: 177 | stdv = 1. / np.sqrt(self.weight.size(1)) 178 | self.bias.data.uniform_(-stdv, stdv) 179 | else: 180 | self.bias_init(self.bias) 181 | 182 | def forward(self, input, adj, *args, **kwargs): 183 | support = T.mm(input, self.weight.t()) 184 | output = T.sparse.mm(adj, support) 185 | if self.bias is not None: 186 | output = output + self.bias 187 | 188 | if self.activation is not None: 189 | output = self.activation(output) 190 | return output 191 | 192 | 193 | class BatchGraphConv(GraphConv): 194 | """ 195 | Performs graph convolution as described in https://arxiv.org/abs/1609.02907 on a batch of graphs. 196 | Adapted from https://github.com/tkipf/pygcn/blob/master/pygcn/layers.py. 197 | 198 | Parameters 199 | ---------- 200 | input_shape 201 | shape of the input tensor. 202 | The first dim should be batch size. 203 | If an integer is passed, it is treated as the size of each input sample. 204 | out_features : int 205 | size of each output sample. 206 | activation 207 | non-linear function to activate the linear result. 208 | It accepts any callable function 209 | as well as a recognizable ``str``. 210 | A list of possible ``str`` is in :func:`~neuralnet_pytorch.utils.function`. 211 | kwargs 212 | extra keyword arguments to pass to activation. 213 | """ 214 | 215 | def __init__(self, input_shape, out_features, bias=True, activation=None, **kwargs): 216 | input_shape = _pointset_shape(input_shape) 217 | super().__init__(input_shape, out_features, bias, activation=activation, **kwargs) 218 | 219 | def forward(self, input, adj, *args, **kwargs): 220 | """ 221 | Performs graphs convolution. 222 | 223 | :param input: 224 | a ``list``/``tuple`` of 2D matrices. 225 | :param adj: 226 | a block diagonal matrix whose diagonal consists of 227 | adjacency matrices of the input graphs. 228 | :return: 229 | a batch of embedded graphs. 230 | """ 231 | 232 | shapes = [input_.shape[0] for input_ in input] 233 | X = T.cat(input, 0) 234 | output = super().forward(X, adj) 235 | output = T.split(output, shapes) 236 | return output 237 | 238 | 239 | @utils.add_custom_repr 240 | class GraphXConv(Module): 241 | """ 242 | Performs GraphX Convolution as described here_. 243 | **Disclaimer:** I am the first author of the paper. 244 | 245 | .. _here: 246 | http://openaccess.thecvf.com/content_ICCV_2019/html/Nguyen_GraphX-Convolution_for_Point_Cloud_Deformation_in_2D-to-3D_Conversion_ICCV_2019_paper.html 247 | 248 | Parameters 249 | ---------- 250 | input_shape 251 | shape of the input tensor. 252 | The first dim should be batch size. 253 | If an integer is passed, it is treated as the size of each input sample. 254 | out_features : int 255 | size of each output sample. 256 | out_instances : int 257 | resolution of the output point clouds. 258 | If not specified, output will have the same resolution as input. 259 | Default: ``None``. 260 | rank 261 | if specified and smaller than `num_out_points`, the mixing matrix will 262 | be broken into a multiplication of two matrices of sizes 263 | ``(num_out_points, rank)`` and ``(rank, input_shape[1])``. 264 | bias 265 | whether to use bias. 266 | Default: ``True``. 267 | activation 268 | non-linear function to activate the linear result. 269 | It accepts any callable function 270 | as well as a recognizable ``str``. 271 | A list of possible ``str`` is in :func:`~neuralnet_pytorch.utils.function`. 272 | weights_init 273 | a kernel initialization method from :mod:`torch.nn.init`. 274 | bias_init 275 | a bias initialization method from :mod:`torch.nn.init`. 276 | kwargs 277 | extra keyword arguments to pass to activation. 278 | """ 279 | 280 | def __init__(self, input_shape, out_features, out_instances=None, rank=None, bias=True, activation=None, 281 | weights_init=None, bias_init=None, **kwargs): 282 | input_shape = _pointset_shape(input_shape) 283 | super().__init__(input_shape=input_shape) 284 | 285 | self.out_features = out_features 286 | self.out_instances = out_instances if out_instances else input_shape[-2] 287 | if rank: 288 | assert rank <= self.out_instances // 2, 'rank should be smaller than half of num_out_points' 289 | 290 | self.rank = rank 291 | self.activation = utils.function(activation, **kwargs) 292 | pattern = list(range(len(input_shape))) 293 | pattern[-1], pattern[-2] = pattern[-2], pattern[-1] 294 | self.pattern = pattern 295 | self.weights_init = weights_init 296 | self.bias_init = bias_init 297 | 298 | self.weight = nn.Parameter(T.Tensor(out_features, input_shape[-1])) 299 | if self.rank is None: 300 | self.mixing = nn.Parameter(T.Tensor(self.out_instances, input_shape[-2])) 301 | else: 302 | self.mixing_u = nn.Parameter(T.Tensor(self.rank, input_shape[-2])) 303 | self.mixing_v = nn.Parameter(T.Tensor(self.out_instances, self.rank)) 304 | 305 | if bias: 306 | self.bias = nn.Parameter(T.Tensor(out_features)) 307 | self.mixing_bias = nn.Parameter(T.Tensor(self.out_instances)) 308 | else: 309 | self.register_parameter('bias', None) 310 | self.register_parameter('mixing_bias', None) 311 | 312 | self.reset_parameters() 313 | 314 | def forward(self, input): 315 | output = utils.dimshuffle(input, self.pattern) 316 | mixing = T.mm(self.mixing_v, self.mixing_u) if self.rank else self.mixing 317 | output = T.matmul(output, mixing.t()) 318 | if self.mixing_bias is not None: 319 | output = output + self.mixing_bias 320 | 321 | output = utils.dimshuffle(output, self.pattern) 322 | output = T.matmul(output, self.weight.t()) 323 | if self.bias is not None: 324 | output = output + self.bias 325 | 326 | return self.activation(output) 327 | 328 | @property 329 | @utils.validate 330 | def output_shape(self): 331 | return self.input_shape[:-2] + (self.out_instances, self.out_features) 332 | 333 | def reset_parameters(self): 334 | super().reset_parameters() 335 | weights_init = partial(nn.init.kaiming_uniform_, a=np.sqrt(5)) if self.weights_init is None \ 336 | else self.weights_init 337 | 338 | weights_init(self.weight) 339 | if self.rank: 340 | weights_init(self.mixing_u) 341 | weights_init(self.mixing_v) 342 | else: 343 | weights_init(self.mixing) 344 | 345 | if self.bias is not None: 346 | if self.bias_init is None: 347 | fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) 348 | bound = 1 / np.sqrt(fan_in) 349 | bias_init = partial(nn.init.uniform_, a=-bound, b=bound) 350 | else: 351 | bias_init = self.bias_init 352 | 353 | bias_init(self.bias) 354 | nn.init.zeros_(self.mixing_bias) 355 | 356 | def extra_repr(self): 357 | s = '{input_shape}, out_features={out_features}, out_instances={out_instances}, rank={rank}' 358 | s = s.format(**self.__dict__) 359 | s += ', activation={}'.format(self.activation.__name__) 360 | return s 361 | --------------------------------------------------------------------------------