MinkowskiEngine多GPU训练

MinkowskiEngine多GPU训练

目前，MinkowskiEngine通过数据并行化支持Multi-GPU训练。在数据并行化中，有一组微型批处理，这些微型批处理将被送到到网络的一组副本中。

首先定义一个网络。

import MinkowskiEngine as ME

from examples.minkunet import MinkUNet34C

 

# Copy the network to GPU

net = MinkUNet34C(3, 20, D=3)

net = net.to(target_device)

同步批处理规范

接下来，创建一个新网络，以ME.MinkowskiSynchBatchNorm替换all ME.MinkowskiBatchNorm。这样一来，网络就可以使用大批处理量，并通过单GPU训练来保持相同的性能。

# Synchronized batch norm

net = ME.MinkowskiSyncBatchNorm.convert_sync_batchnorm(net);

接下来，需要创建网络和最终损耗层的副本（如果使用一个副本）。

import torch.nn.parallel as parallel

 

criterion = nn.CrossEntropyLoss()

criterions = parallel.replicate(criterion, devices)

加载多个批次

在训练过程中，每次训练迭代都需要一组微型批次。使用了一个返回一个mini-batches批处理的函数，但是无需遵循这种模式。

# Get new data

inputs, labels = [], []

for i in range(num_devices):

    coords, feat, label = data_loader()  // parallel data loaders can be used

    with torch.cuda.device(devices[i]):

      inputs.append(ME.SparseTensor(feat, coords=coords).to(devices[i]))

    labels.append(label.to(devices[i]))

将weights复制到设备

首先，将权重复制到所有设备。

replicas = parallel.replicate(net, devices)

将副本应用于所有批次

接下来，将所有mini-batches批次送到到所有设备上网络的相应副本。然后将所有输出要素输入损耗层。

outputs = parallel.parallel_apply(replicas, inputs, devices=devices)

 

# Extract features from the sparse tensors to use a pytorch criterion

out_features = [output.F for output in outputs]

losses = parallel.parallel_apply(

    criterions, tuple(zip(out_features, labels)), devices=devices)

收集所有损失到目标设备。

loss = parallel.gather(losses, target_device, dim=0).mean()

其余训练（如backward训练和在优化器中采取向前步骤）类似于单GPU训练。请参阅完整的multi-gpu示例以获取更多详细信息。

import os
	import argparse
	import numpy as np
	from time import time
	from urllib.request import urlretrieve
	try:
	import open3d as o3d
	except ImportError:
	raise ImportError("Please install open3d-python with `pip install open3d`.")
	import torch
	import torch.nn as nn
	from torch.optim import SGD
	import MinkowskiEngine as ME
	from examples.minkunet import MinkUNet34C
	import torch.nn.parallel as parallel
	if not os.path.isfile("weights.pth"):
	urlretrieve("http://cvgl.stanford.edu/data2/minkowskiengine/1.ply", "1.ply")
	parser = argparse.ArgumentParser()
	parser.add_argument("--file_name", type=str, default="1.ply")
	parser.add_argument("--batch_size", type=int, default=4)
	parser.add_argument("--max_ngpu", type=int, default=2)
	cache = {}
	def load_file(file_name, voxel_size):
	if file_name not in cache:
	pcd = o3d.io.read_point_cloud(file_name)
	cache[file_name] = pcd
	pcd = cache[file_name]
	quantized_coords, feats = ME.utils.sparse_quantize(
	np.array(pcd.points, dtype=np.float32),
	np.array(pcd.colors, dtype=np.float32),
	quantization_size=voxel_size,
	)
	random_labels = torch.zeros(len(feats))
	return quantized_coords, feats, random_labels
	def generate_input(file_name, voxel_size):
	# Create a batch, this process is done in a data loader during training in parallel.
	batch = [load_file(file_name, voxel_size)]
	coordinates_, featrues_, labels_ = list(zip(*batch))
	coordinates, features, labels = ME.utils.sparse_collate(
	coordinates_, featrues_, labels_
	)
	# Normalize features and create a sparse tensor
	return coordinates, (features - 0.5).float(), labels
	if __name__ == "__main__":
	# loss and network
	config = parser.parse_args()
	num_devices = torch.cuda.device_count()
	num_devices = min(config.max_ngpu, num_devices)
	devices = list(range(num_devices))
	print("''''''''''''''''''''''''''''''''''''''''''''''''''''''''''")
	print("' WARNING: This example is deprecated. '")
	print("' Please use DistributedDataParallel or pytorch-lightning'")
	print("''''''''''''''''''''''''''''''''''''''''''''''''''''''''''")
	print(
	f"Testing {num_devices} GPUs. Total batch size: {num_devices * config.batch_size}"
	)
	# For copying the final loss back to one GPU
	target_device = devices[0]
	# Copy the network to GPU
	net = MinkUNet34C(3, 20, D=3)
	net = net.to(target_device)
	# Synchronized batch norm
	net = ME.MinkowskiSyncBatchNorm.convert_sync_batchnorm(net)
	optimizer = SGD(net.parameters(), lr=1e-1)
	# Copy the loss layer
	criterion = nn.CrossEntropyLoss()
	criterions = parallel.replicate(criterion, devices)
	min_time = np.inf
	for iteration in range(10):
	optimizer.zero_grad()
	# Get new data
	inputs, all_labels = [], []
	for i in range(num_devices):
	coordinates, features, labels = generate_input(config.file_name, 0.05)
	with torch.cuda.device(devices[i]):
	inputs.append(ME.SparseTensor(features, coordinates, device=devices[i]))
	all_labels.append(labels.long().to(devices[i]))
	# The raw version of the parallel_apply
	st = time()
	replicas = parallel.replicate(net, devices)
	outputs = parallel.parallel_apply(replicas, inputs, devices=devices)
	# Extract features from the sparse tensors to use a pytorch criterion
	out_features = [output.F for output in outputs]
	losses = parallel.parallel_apply(
	criterions, tuple(zip(out_features, all_labels)), devices=devices
	)
	loss = parallel.gather(losses, target_device, dim=0).mean()
	# Gradient
	loss.backward()
	optimizer.step()
	t = time() - st
	min_time = min(t, min_time)
	print(
	f"Iteration: {iteration}, Loss: {loss.item()}, Time: {t}, Min time: {min_time}"
	)
	# Must clear cache at regular interval
	if iteration % 10 == 0:
	torch.cuda.empty_cache()

 

加速实验

在4x Titan XP上使用各种批次大小进行实验，并将负载平均分配给每个GPU。例如，使用1个GPU，每个批次将具有8个批处理大小。使用2个GPU，每个GPU将具有4个批次。使用4个GPU，每个GPU的批处理大小为2。

GPU数量	每个GPU的批量大小	每次迭代时间	加速（理想）
1个GPU	8	1.611秒	x1（x1）
2个GPU	4	0.916秒	x1.76（x2）
4个GPU	2	0.689秒	x2.34（x4）
GPU数量	每个GPU的批量大小	每次迭代时间	加速（理想）
1个GPU	12	2.691秒	x1（x1）
2个GPU	6	1.413秒	x1.90（x2）
3个GPU	4	1.064秒	x2.53（x3）
4个GPU	3	1.006秒	x2.67（x4）

 

GPU数量	每个GPU的批量大小	每次迭代时间	加速（理想）
1个GPU	16	3.543秒	x1（x1）
2个GPU	8	1.933秒	x1.83（x2）
4个GPU	4	1.322秒	x2.68（x4）
GPU数量	每个GPU的批量大小	每次迭代时间	加速（理想）
1个GPU	18岁	4.391秒	x1（x1）
2个GPU	9	2.114秒	x2.08（x2）
3个GPU	6	1.660秒	x2.65（x3）

 

GPU数量	每个GPU的批量大小	每次迭代时间	加速（理想）
1个GPU	20	4.639秒	x1（x1）
2个GPU	10	2.426秒	x1.91（x2）
4个GPU	5	1.707秒	x2.72（x4）
GPU数量	每个GPU的批量大小	每次迭代时间	加速（理想）
1个GPU	21	4.894秒	x1（x1）
3个GPU	7	1.877秒	x2.61（x3）

分析

批量较小时，加速非常适中。对于大批处理大小（例如18和20），随着线程初始化开销在大工作量上摊销，速度会提高。

同样，在所有情况下，使用4个GPU效率都不高，并且速度似乎很小（总批量大小为18的3-GPU的x2.65与总批量大小为20的4-GPU的x2.72）。因此，建议最多使用3个大批量的GPU。

GPU数量	平均加速（理想）
1个GPU	x1（x1）
2个GPU	x1.90（x2）
3个GPU	x2.60（x3）
4个GPU	x2.60（x4）

适度加速的原因是由于CPU使用率过高。在Minkowski引擎中，所有稀疏张量坐标都在CPU上进行管理，并且内核in-out出入图需要大量的CPU计算。因此，为了提高速度，建议使用更快的CPU，这可能是大点云的瓶颈。