训练流程¶

Minkowski Engine 可以与 PyTorch torch.utils.data.DataLoader 无缝协作。在继续之前，请确保您熟悉数据加载教程 torch.utils.data.DataLoader。

制作数据集¶

您需要做的第一件事是加载或生成数据。如果您使用自己的数据集，这将是最耗时的部分。但是，这是繁琐的，而不是困难的:) 如果您继承 torch.utils.data.Dataset，您需要填写的最重要的部分是 __getitem__(self, index)；如果您继承 torch.utils.data.IterableDataset，则需要填写 __iter__(self)。在此示例中，让我们创建一个生成噪声线的随机数据集。

class RandomLineDataset(torch.utils.data.Dataset):

    ...

    def __getitem__(self, i):
        # Regardless of the input index, return randomized data
        angle, intercept = np.tan(self._uniform_to_angle(
            self.rng.rand())), self.rng.rand()

        # Line as x = cos(theta) * t, y = sin(theta) * t + intercept and random t's
        # Drop some samples
        xs_data = self._sample_xs(self.num_data)
        ys_data = angle * xs_data + intercept + self._sample_noise(
            self.num_data, [0, 0.1])

        noise = 4 * (self.rng.rand(self.num_noise, 2) - 0.5)

        # Concatenate data
        input = np.vstack([np.hstack([xs_data, ys_data]), noise])
        feats = input
        labels = np.vstack(
            [np.ones((self.num_data, 1)),
             np.zeros((self.num_noise, 1))]).astype(np.int32)

        ...

        # quantization step
        return various_outputs

在这里，我创建了一个继承 torch.utils.data.Dataset 的数据集，但是您可以继承 torch.utils.data.IterableDataset 并填写 __iter__(self) 来代替。

量化¶

我们使用稀疏张量作为输入。像任何张量一样，稀疏张量值在离散位置（索引）处定义。因此，量化定义了特征的坐标是关键步骤，并且 quantization_size 是一个重要的超参数，它会极大地影响网络的性能。您必须选择正确的量化大小以及正确地量化坐标。

Minkowski Engine 提供了一组用于量化和稀疏张量生成的快速且优化的函数。在这里，我们使用 MinkowskiEngine.utils.sparse_quantize。

class RandomLineDataset(torch.utils.data.Dataset):

    ...

    def __getitem__(self, i):

        ...

        # Quantize the input
        discrete_coords, unique_feats, unique_labels = ME.utils.sparse_quantize(
            coords=input,
            feats=feats,
            labels=labels,
            quantization_size=self.quantization_size)

        return discrete_coords, unique_feats, unique_labels

量化坐标的另一种方法是使用返回的映射索引。如果您有非常规的输入，这将很有用。

class RandomLineDataset(torch.utils.data.Dataset):

    ...

    def __getitem__(self, i):

        ...

        coords /= self.quantization_size

        # Quantize the input
        mapping = ME.utils.sparse_quantize(
            coords=coords,
            return_index=True)

        return coords[mapping], feats[mapping], labels[mapping]

制作 DataLoader¶

创建数据集后，您需要一个数据加载器来调用数据集并为神经网络训练生成小批量。这部分相对容易，但是我们必须使用自定义的 collate_fn 来生成合适的稀疏张量。

train_dataset = RandomLineDataset(...)
# Option 1
train_dataloader = DataLoader(
    train_dataset,
    ...
    collate_fn=ME.utils.SparseCollation())

# Option 2
train_dataloader = DataLoader(
    train_dataset,
    ...
    collate_fn=ME.utils.batch_sparse_collate)

在这里，我们可以使用提供的整理类 MinkowskiEngine.utils.SparseCollation 或函数 MinkowskiEngine.utils.batch_sparse_collate 将输入转换为适当的输出，我们可以使用该输出初始化稀疏张量。但是，如果您需要自己的整理函数，则可以按照以下示例进行操作。

def custom_collation_fn(data_labels):
    coords, feats, labels = list(zip(*data_labels))

    # Create batched coordinates for the SparseTensor input
    bcoords = ME.utils.batched_coordinates(coords)

    # Concatenate all lists
    feats_batch = torch.from_numpy(np.concatenate(feats, 0)).float()
    labels_batch = torch.from_numpy(np.concatenate(labels, 0)).int()

    return bcoords, feats, labels

...

train_dataset = RandomLineDataset(...)
train_dataloader = DataLoader(
    train_dataset,
    ...
    collate_fn=custom_collation_fn)

训练¶

有了所有东西之后，让我们创建一个网络并使用生成的数据对其进行训练。需要注意的一件事是，如果为数据加载器使用多个 num_workers，则必须确保 MinkowskiEngine.SparseTensor 生成部分必须位于主 python 进程中，因为所有 python 多进程都使用单独的进程，并且 MinkowskiEngine.CoordsManager（维护坐标哈希表和内核映射的内部 C++ 结构）无法在生成它的进程外部引用。

# Binary classification
net = UNet(
    2,  # in nchannel
    2,  # out_nchannel
    D=2)

optimizer = optim.SGD(
    net.parameters(),
    lr=config.lr,
    momentum=config.momentum,
    weight_decay=config.weight_decay)

criterion = torch.nn.CrossEntropyLoss(ignore_index=-100)

# Dataset, data loader
train_dataset = RandomLineDataset(noise_type='gaussian')

train_dataloader = DataLoader(
    train_dataset,
    batch_size=config.batch_size,
    collate_fn=collation_fn,
    num_workers=1)

for epoch in range(config.max_epochs):
    train_iter = iter(train_dataloader)

    # Training
    net.train()
    for i, data in enumerate(train_iter):
        coords, feats, labels = data
        out = net(ME.SparseTensor(feats, coords))
        optimizer.zero_grad()
        loss = criterion(out.F.squeeze(), labels.long())
        loss.backward()
        optimizer.step()

        accum_loss += loss.item()
        accum_iter += 1
        tot_iter += 1

        if tot_iter % 10 == 0 or tot_iter == 1:
            print(
                f'Iter: {tot_iter}, Epoch: {epoch}, Loss: {accum_loss / accum_iter}'
            )
            accum_loss, accum_iter = 0, 0

最后，一旦您组装了所有代码，就可以训练您的网络了。

$ python -m examples.training

Epoch: 0 iter: 1, Loss: 0.7992178201675415
Epoch: 0 iter: 10, Loss: 0.5555745628145006
Epoch: 0 iter: 20, Loss: 0.4025680094957352
Epoch: 0 iter: 30, Loss: 0.3157463788986206
Epoch: 0 iter: 40, Loss: 0.27348957359790804
Epoch: 0 iter: 50, Loss: 0.2690591633319855
Epoch: 0 iter: 60, Loss: 0.258208692073822
Epoch: 0 iter: 70, Loss: 0.34842072874307634
Epoch: 0 iter: 80, Loss: 0.27565130293369294
Epoch: 0 iter: 90, Loss: 0.2860450878739357
Epoch: 0 iter: 100, Loss: 0.24737665355205535
Epoch: 1 iter: 110, Loss: 0.2428090125322342
Epoch: 1 iter: 120, Loss: 0.25397603064775465
Epoch: 1 iter: 130, Loss: 0.23624965399503708
Epoch: 1 iter: 140, Loss: 0.2247777447104454
Epoch: 1 iter: 150, Loss: 0.22956613600254058
Epoch: 1 iter: 160, Loss: 0.22803852707147598
Epoch: 1 iter: 170, Loss: 0.24081039279699326
Epoch: 1 iter: 180, Loss: 0.22322929948568343
Epoch: 1 iter: 190, Loss: 0.22531934976577758
Epoch: 1 iter: 200, Loss: 0.2116936132311821
...

可以在 examples/training.py 中找到原始代码。