def setup_multi_device(net):
from CaffeLoader import ModelParallel
device_splits = params.multidevice_strategy.split(',')
device_list = []
for i, device in enumerate(params.gpu):
if str(device).lower() != 'c':
device_list.append("cuda:" + str(device))
else:
device_list.append("cpu")
cur_chunk = nn.Sequential()
chunks = []
for i, l in enumerate(net):
cur_chunk.add_module(str(i), net[i])
if str(i) in device_splits and device_splits != '':
del device_splits[0]
chunks.append(cur_chunk)
cur_chunk = nn.Sequential()
chunks.append(cur_chunk)
for i, chunk in enumerate(chunks):
chunk.to(device_list[i])
new_net = ModelParallel(chunks, device_list)
return new_net
def setup_multi_device(net):
assert len(params.gpu.split(',')) - 1 == len(params.multidevice_strategy.split(',')), \
"The number of -multidevice_strategy layer indices minus 1, must be equal to the number of -gpu devices."
new_net = ModelParallel(net, params.gpu, params.multidevice_strategy)
return new_net
def __setup_multi_device(self, gpu, multidevice_strategy):
from CaffeLoader import ModelParallel
self.multidevice = True
assert len(gpu.split(',')) - 1 == len(multidevice_strategy.split(',')), \
"The number of -multidevice_strategy layer indices minus 1, must be equal to the number of -gpu devices."
self.net = ModelParallel(self.net, gpu, multidevice_strategy)