闵可夫斯基引擎Minkowski Engine

闵可夫斯基引擎Minkowski Engine

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

pip install git+https://github.com/NVIDIA/MinkowskiEngine.git

 

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

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

但是,在这项工作中,专注于稀疏的数据,尤其是空间稀疏的高维输入。还可以将这些数据表示为稀疏张量,并且这些稀疏张量在3D感知,配准和统计数据等高维问题中很常见。将专门用于这些输入的神经网络定义为稀疏张量网络,这些稀疏张量网络处理并生成稀疏张量作为输出。为了构建稀疏张量网络,建立了所有标准的神经网络层,例如MLP,非线性,卷积,规范化,池化操作,就像在密集张量上定义,并在Minkowski引擎中实现的方法一样。

在下面的稀疏张量卷积上可视化了一个稀疏张量网络操作。稀疏张量上的卷积层与密集张量上的卷积层相似。但是,在稀疏张量上,在一些指定点上计算卷积输出,这些点可以在广义卷积中进行控制。

特征

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

Requirements

  • Ubuntu >= 14.04
  • 11.1 > CUDA >= 10.1.243
  • pytorch >= 1.5
  • python >= 3.6
  • GCC >= 7

Pip

MinkowskiEngine是通过PyPI MinkowskiEngine分发的,可以使用简单安装pip。按照说明安装pytorch 。接下来,安装openblas

sudo apt install libopenblas-dev
pip install torch
pip install -U MinkowskiEngine --install-option="--blas=openblas" -v
 
# For pip installation from the latest source
# pip install -U git+https://github.com/NVIDIA/MinkowskiEngine

If you want to specify arguments for the setup script, please refer to the following command.

# Uncomment some options if things don't work
pip install -U git+https://github.com/NVIDIA/MinkowskiEngine \
#                           \ # 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 when torch fails to find cuda_home.
#                           --install-option="--cuda_home=/usr/local/cuda" \
#                           \ # uncomment the following line to override to openblas, atlas, mkl, blas
#                           --install-option="--blas=openblas" \

快速启动

要使用Minkowski引擎,首先需要导入引擎。然后,将需要定义网络。如果没有量化数据,则需要将(空间)数据体素化或量化为稀疏张量。幸运的是,Minkowski引擎提供了量化功能(MinkowskiEngine.utils.sparse_quantize)。

Anaconda

We recommend python>=3.6 for installation. First, follow the anaconda documentation to install anaconda on your computer.

sudo apt install libopenblas-dev
conda create -n py3-mink python=3.8
conda activate py3-mink
conda install numpy mkl-include pytorch cudatoolkit=11.0 -c pytorch
pip install -U git+https://github.com/NVIDIA/MinkowskiEngine

System Python

Like the anaconda installation, make sure that you install pytorch with the same CUDA version that nvcc uses.

# install system requirements
sudo apt install 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
 
git clone https://github.com/NVIDIA/MinkowskiEngine.git
 
cd MinkowskiEngine
 
python setup.py install
# To specify blas, CUDA_HOME and force CUDA installation, use the following command
# python setup.py install --blas=openblas --cuda_home=/usr/local/cuda --force_cuda

Creating a Network

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,
                has_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)

Forward and backward using the custom network

    # 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, coords=coords)
    # Forward
    output = net(input)
 
    # Loss
    loss = criterion(output.F, label)

 

posted @ 2021-01-03 13:55  吴建明wujianming  阅读(2304)  评论(0编辑  收藏  举报