使用 Pytorch 层

MinkowskiEngine.SparseTensor 是 torch.Tensor 的一个浅封装。因此，将稀疏张量转换为 pytorch 张量非常容易，反之亦然。

示例：用于分类的特征

在此示例中，我们将展示如何从 MinkowskiEngine.SparseTensor 中提取特征，并将这些特征与 pytorch 层一起使用。

首先，让我们创建一个网络，为最小批次中的每个输入生成一个特征向量。

import torch.nn as nn
import MinkowskiEngine as ME


class ExampleNetwork(nn.Module):

    def __init__(self, in_feat, out_feat, D):
        self.net = 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(),
            ME.MinkowskiConvolution(
                in_channels=64,
                out_channels=128,
                kernel_size=3,
                stride=2,
                dimension=D), ME.MinkowskiBatchNorm(128), ME.MinkowskiReLU(),
            ME.MinkowskiGlobalPooling(),
            ME.MinkowskiLinear(128, out_feat))

    def forward(self, x):
        return self.net(x)

请注意，上面的 MinkowskiEngine.MinkowskiGlobalPooling 层会对输入稀疏张量中的所有特征进行平均，并在 \(B\) 是批次大小（自适应改变）且 \(D_F\) 是输入稀疏张量的特征维度时，生成 \(B \times D_F\)。

然后，在训练期间，我们可以通过访问稀疏张量的特征 MinkowskiEngine.SparseTensor.F 或 MinkowskiEngine.SparseTensor.feats 来使用 torch.nn.CrossEntropyLoss 层。

criterion = nn.CrossEntropyLoss()

for i in range(10):
    optimizer.zero_grad()

    # Get new data
    coords, feat, label = data_loader()
    input = ME.SparseTensor(features=feat, coordinates=coords, device=device)
    label = label.to(device)

    # Forward
    output = net(input)

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

有关完整的演示，请参考 examples/example.py。