MinkowskiEngine:基于PyTorch的稀疏张量自动微分库项目

Minkowski Engine is an auto-diff neural network library for high-dimensional sparse tensors

分支3Tags18
文件最后提交记录最后更新时间
5 年前
3 年前
3 年前
5 年前
4 年前
4 年前
3 年前
4 年前
6 年前
4 年前
6 年前
6 年前
5 年前
3 年前
5 年前
3 年前

Minkowski Engine

PyPI 版本 PyPI 每月下载 Slack 聊天

Minkowski Engine 是一个用于稀疏张量的自动微分库。它支持所有标准神经网络层,如卷积、池化、反池化和广播操作,适用于稀疏张量。更多信息,请访问 官方文档页面

新闻

  • 2021-08-11 添加了 Docker 安装说明
  • 2021-08-06 已解决与 pytorch 1.8 和 1.9 相关的所有安装错误。
  • 2021-04-08 由于最近的 pytorch 1.8 + CUDA 11 问题,建议使用 anaconda 进行安装
  • 2020-12-24 v0.5 版本现已可用!新版本为所有坐标管理函数提供了 CUDA 加速。

示例网络

Minkowski Engine 支持在稀疏张量上构建的各种功能。我们在下面列出了一些流行的网络架构和应用。要运行示例,请安装该包,并在包的根目录下运行命令。

示例 网络和命令
语义分割

python -m examples.indoor
分类
python -m examples.classification_modelnet40
重建

python -m examples.reconstruction
完成度
python -m examples.completion
检测

稀疏张量网络:面向空间稀疏张量的神经网络

压缩神经网络以加速推理和最小化内存占用已被广泛研究。模型压缩的一种流行技术是修剪卷积网络中的权重,也称为稀疏卷积网络。这种参数空间稀疏性用于模型压缩,压缩了在密集张量上运行且所有中间激活也为密集张量的网络。

然而,在这项工作中,我们关注的是空间稀疏数据,尤其是高维输入和三维数据以及三维对象表面的卷积,最初在 Siggraph'17 中提出。我们可以将这些数据表示为稀疏张量,这些稀疏张量在高维问题(如三维感知、注册和统计数据)中很常见。我们定义了专门用于这些输入的神经网络,称为 稀疏张量网络,并且这些稀疏张量网络处理和生成稀疏张量作为输出。为了构建稀疏张量网络,我们按照在密集张量上定义和实现的同样方式构建了所有标准神经网络层,如 MLPs、非线性、卷积、归一化、池化操作。

我们下面可视化了一个稀疏张量网络操作,卷积。在稀疏张量上的卷积层与在密集张量上的工作方式相似。但是,在稀疏张量上,我们在少数指定的点上计算卷积输出,这些点可以在 广义卷积 中控制。更多信息,请访问 稀疏张量网络的文档页面术语页面

密集张量 稀疏张量

特点

  • 无限高维稀疏张量支持
  • 所有标准神经网络层(卷积、池化、广播等)
  • 动态计算图
  • 自定义核形状
  • 多 GPU 训练
  • 多线程内核映射
  • 多线程编译
  • 高度优化的 GPU 内核

环境要求

  • Ubuntu >= 14.04
  • CUDA >= 10.1.243 且 与 pytorch 使用的 CUDA 版本相同(例如:如果使用 conda cudatoolkit=11.1,则安装 MinkowskiEngine 时也使用 CUDA=11.1)
  • pytorch >= 1.7 指定 CUDA 版本时,请使用 conda 进行安装。您必须匹配 pytorch 使用的 CUDA 版本和 Minkowski Engine 安装时使用的 CUDA 版本。conda install -y -c nvidia -c pytorch pytorch=1.8.1 cudatoolkit=10.2
  • python >= 3.6
  • ninja(用于安装)
  • GCC >= 7.4.0

安装

您可以使用 pip、anaconda 或直接在系统上安装 Minkowski Engine。如果在安装包时遇到问题,请查看 安装维基页面。 如果找不到相关的问题,请在 GitHub 问题页面 上报告问题。

Pip

MinkowskiEngine 通过 PyPI MinkowskiEngine 进行分发,您可以使用 pip 命令轻松安装。 首先,请按照 说明 安装 pytorch。接下来,安装 openblas

sudo apt install build-essential python3-dev libopenblas-dev
pip install torch ninja
pip install -U MinkowskiEngine --install-option="--blas=openblas" -v --no-deps

# For pip installation from the latest source
# pip install -U git+https://github.com/NVIDIA/MinkowskiEngine --no-deps

如果您需要为设置脚本指定参数,请参考以下命令。

# Uncomment some options if things don't work
# export CXX=c++; # set this if you want to use a different C++ compiler
# export CUDA_HOME=/usr/local/cuda-11.1; # or select the correct cuda version on your system.
pip install -U git+https://github.com/NVIDIA/MinkowskiEngine -v --no-deps \
#                           \ # uncomment the following line if you want to force cuda installation
#                           --install-option="--force_cuda" \
#                           \ # uncomment the following line if you want to force no cuda installation. force_cuda supercedes cpu_only
#                           --install-option="--cpu_only" \
#                           \ # uncomment the following line to override to openblas, atlas, mkl, blas
#                           --install-option="--blas=openblas" \

Anaconda

MinkowskiEngine 支持CUDA 10.2与CUDA 11.1,这两种版本适用于大多数最新的 PyTorch 版本。

CUDA 10.2

我们建议安装时使用 python>=3.6。 首先,请按照 Anaconda 官方文档 的指引,在您的计算机上安装 Anaconda。

sudo apt install g++-7  # For CUDA 10.2, must use GCC < 8
# Make sure `g++-7 --version` is at least 7.4.0
conda create -n py3-mink python=3.8
conda activate py3-mink

conda install openblas-devel -c anaconda
conda install pytorch=1.9.0 torchvision cudatoolkit=10.2 -c pytorch -c nvidia

# Install MinkowskiEngine
export CXX=g++-7
# Uncomment the following line to specify the cuda home. Make sure `$CUDA_HOME/nvcc --version` is 10.2
# export CUDA_HOME=/usr/local/cuda-10.2
pip install -U git+https://github.com/NVIDIA/MinkowskiEngine -v --no-deps --install-option="--blas_include_dirs=${CONDA_PREFIX}/include" --install-option="--blas=openblas"

# Or if you want local MinkowskiEngine
git clone https://github.com/NVIDIA/MinkowskiEngine.git
cd MinkowskiEngine
export CXX=g++-7
python setup.py install --blas_include_dirs=${CONDA_PREFIX}/include --blas=openblas

CUDA 11.X

我们建议安装时使用 python>=3.6 版本。 首先,请遵循 Anaconda 官方文档 在您的计算机上安装 Anaconda。

conda create -n py3-mink python=3.8
conda activate py3-mink

conda install openblas-devel -c anaconda
conda install pytorch=1.9.0 torchvision cudatoolkit=11.1 -c pytorch -c nvidia

# Install MinkowskiEngine

# Uncomment the following line to specify the cuda home. Make sure `$CUDA_HOME/nvcc --version` is 11.X
# export CUDA_HOME=/usr/local/cuda-11.1
pip install -U git+https://github.com/NVIDIA/MinkowskiEngine -v --no-deps --install-option="--blas_include_dirs=${CONDA_PREFIX}/include" --install-option="--blas=openblas"

# Or if you want local MinkowskiEngine
git clone https://github.com/NVIDIA/MinkowskiEngine.git
cd MinkowskiEngine
python setup.py install --blas_include_dirs=${CONDA_PREFIX}/include --blas=openblas

系统Python环境

与Anaconda安装类似,确保安装与nvcc所使用的相同版本的CUDA的PyTorch。

# install system requirements
sudo apt install build-essential python3-dev libopenblas-dev

# Skip if you already have pip installed on your python3
curl https://bootstrap.pypa.io/get-pip.py | python3

# Get pip and install python requirements
python3 -m pip install torch numpy ninja

git clone https://github.com/NVIDIA/MinkowskiEngine.git

cd MinkowskiEngine

python setup.py install
# To specify blas, CXX, CUDA_HOME and force CUDA installation, use the following command
# export CXX=c++; export CUDA_HOME=/usr/local/cuda-11.1; python setup.py install --blas=openblas --force_cuda

容器化技术——Docker

在当今的软件开发与运维领域,Docker 作为一种革命性的容器化技术,正日益受到广泛关注。它通过将应用及其依赖环境打包,实现了应用的快速部署和高效运行。Docker 的高效性、灵活性与可移植性,为开发者和运维人员提供了一种全新的工作模式,大幅缩短了从开发到上线的时间周期,优化了资源的使用效率,确保了应用在不同环境下的稳定运行。

git clone https://github.com/NVIDIA/MinkowskiEngine
cd MinkowskiEngine
docker build -t minkowski_engine docker

一旦 Docker 镜像构建完成,请检查它是否正确加载了 MinkowskiEngine。

docker run MinkowskiEngine python3 -c "import MinkowskiEngine; print(MinkowskiEngine.__version__)"

仅CPU构建与BLAS配置(MKL)

Minkowski Engine支持在其他不具备NVidia GPU的平台进行仅CPU构建。请参考快速入门以获取更多详细信息。

快速入门

要使用Minkowski Engine,首先需要导入该引擎。 接着,你需要定义网络。如果你的数据未量化,则需要将(空间)数据体素化或量化为稀疏张量。 幸运的是,Minkowski Engine提供了量化函数(MinkowskiEngine.utils.sparse_quantize)。

创建网络

import torch.nn as nn
import MinkowskiEngine as ME

class ExampleNetwork(ME.MinkowskiNetwork):

    def __init__(self, in_feat, out_feat, D):
        super(ExampleNetwork, self).__init__(D)
        self.conv1 = nn.Sequential(
            ME.MinkowskiConvolution(
                in_channels=in_feat,
                out_channels=64,
                kernel_size=3,
                stride=2,
                dilation=1,
                bias=False,
                dimension=D),
            ME.MinkowskiBatchNorm(64),
            ME.MinkowskiReLU())
        self.conv2 = nn.Sequential(
            ME.MinkowskiConvolution(
                in_channels=64,
                out_channels=128,
                kernel_size=3,
                stride=2,
                dimension=D),
            ME.MinkowskiBatchNorm(128),
            ME.MinkowskiReLU())
        self.pooling = ME.MinkowskiGlobalPooling()
        self.linear = ME.MinkowskiLinear(128, out_feat)

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.pooling(out)
        return self.linear(out)

使用自定义网络进行前后向传播

    # loss and network
    criterion = nn.CrossEntropyLoss()
    net = ExampleNetwork(in_feat=3, out_feat=5, D=2)
    print(net)

    # a data loader must return a tuple of coords, features, and labels.
    coords, feat, label = data_loader()
    input = ME.SparseTensor(feat, coordinates=coords)
    # Forward
    output = net(input)

    # Loss
    loss = criterion(output.F, label)

讨论与文档

如有讨论或疑问,请使用 minkowskiengine@googlegroups.com 联系我们。 有关 API 和一般性使用,请参考 MinkowskiEngine 文档页面 获取更多详细信息。

对于 API 中未列出的问题以及功能请求,欢迎在 GitHub 问题页面 提交问题。

已知问题

指定 CUDA 架构列表

在某些情况下,您需要显式指定您的 GPU 使用哪种计算能力。默认列表可能不包含您的架构。

export TORCH_CUDA_ARCH_LIST="5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX"; python setup.py install --force_cuda

未处理的内存不足 thrust::system 异常

在 CUDA 10 版本中存在一个已知的 thrust 问题,会导致未处理的 thrust 异常。请参考 问题详情

过多的 GPU 内存使用或频繁内存不足

这个错误可能有几个原因。

  1. 在长时间运行的训练过程中内存不足

MinkowskiEngine 是一个专用库,可以在每次迭代中处理不同数量的点或不同数量的非零元素,这在点云数据中很常见。 然而,pytorch 的实现是假设在每次迭代中点的数量或激活的大小不会改变。因此,pytorch 使用的 GPU 内存缓存可能导致不必要的内存消耗。

具体来说,pytorch 缓存了内存空间的块,以加速在每个张量创建时使用的分配。如果它找不到内存空间,它会拆分现有的缓存内存,或者在没有任何足够大的缓存内存时分配新的空间。因此,每次我们在 pytorch 中使用不同数量的点(非零元素的数量)时,它要么拆分现有缓存,要么预留新的内存。如果缓存过于碎片化且占用了所有 GPU 空间,它将引发内存不足错误。

为防止这种情况,您必须定期使用 torch.cuda.empty_cache() 清除缓存。

CUDA 11.1 安装

wget https://developer.download.nvidia.com/compute/cuda/11.1.1/local_installers/cuda_11.1.1_455.32.00_linux.run
sudo sh cuda_11.1.1_455.32.00_linux.run --toolkit --silent --override

# Install MinkowskiEngine with CUDA 11.1
export CUDA_HOME=/usr/local/cuda-11.1; pip install MinkowskiEngine -v --no-deps

在具有大量CPU节点的设备上运行MinkowskiEngine

MinkowskiEngine采用OpenMP来并行化内核映射生成。然而,当用于并行的线程数过大(例如 OMP_NUM_THREADS=80)时,由于所有线程都在等待多线程锁的释放,效率会迅速下降。 在这种情况下,请设置用于OpenMP的线程数。通常,低于24的任何数值都是可行的,但请在您的系统上寻找最佳配置。

export OMP_NUM_THREADS=<number of threads to use>; python <your_program.py>

引用 Minkowski Engine

如果您使用了 Minkowski Engine,请引用以下文献:

@inproceedings{choy20194d,
  title={4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural Networks},
  author={Choy, Christopher and Gwak, JunYoung and Savarese, Silvio},
  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
  pages={3075--3084},
  year={2019}
}

对于多线程内核映射生成,请引用如下:

@inproceedings{choy2019fully,
  title={Fully Convolutional Geometric Features},
  author={Choy, Christopher and Park, Jaesik and Koltun, Vladlen},
  booktitle={Proceedings of the IEEE International Conference on Computer Vision},
  pages={8958--8966},
  year={2019}
}

对于高维卷积的步进池化层,请引用以下文献:

@inproceedings{choy2020high,
  title={High-dimensional Convolutional Networks for Geometric Pattern Recognition},
  author={Choy, Christopher and Lee, Junha and Ranftl, Rene and Park, Jaesik and Koltun, Vladlen},
  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
  year={2020}
}

对于生成转置卷积,请引用如下:

@inproceedings{gwak2020gsdn,
  title={Generative Sparse Detection Networks for 3D Single-shot Object Detection},
  author={Gwak, JunYoung and Choy, Christopher B and Savarese, Silvio},
  booktitle={European conference on computer vision},
  year={2020}
}

单元测试

进行单元测试和梯度检查时,请使用 torch 1.7 或更高版本。

使用 Minkowski Engine 的项目

请随时更新维基页面,添加您的项目!