in dlrm_s_caffe2.py [0:0]
def AddLayerWrapper(self, layer, inp_blobs, out_blobs,
add_prefix=True, reset_grad=False, **kwargs):
# auxiliary routine to adjust tags
def adjust_tag(blobs, on_device):
if blobs.__class__ == str:
_blobs = on_device + blobs
elif blobs.__class__ == list:
_blobs = list(map(lambda tag: on_device + tag, blobs))
else: # blobs.__class__ == model_helper.ModelHelper or something else
_blobs = blobs
return _blobs
if self.ndevices > 1 and add_prefix:
# add layer on multiple devices
ll = []
for d in range(self.ndevices):
# add prefix on_device
on_device = "gpu_" + str(d) + "/"
_inp_blobs = adjust_tag(inp_blobs, on_device)
_out_blobs = adjust_tag(out_blobs, on_device)
# WARNING: reset_grad option was exlusively designed for WeightedSum
# with inp_blobs=[w, tag_one, "", lr], where "" will be replaced
if reset_grad:
w_grad = self.gradientMap[_inp_blobs[0]]
_inp_blobs[2] = w_grad
# add layer to the model
with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, d)):
if kwargs:
new_layer = layer(_inp_blobs, _out_blobs, **kwargs)
else:
new_layer = layer(_inp_blobs, _out_blobs)
ll.append(new_layer)
return ll
else:
# add layer on a single device
# WARNING: reset_grad option was exlusively designed for WeightedSum
# with inp_blobs=[w, tag_one, "", lr], where "" will be replaced
if reset_grad:
w_grad = self.gradientMap[inp_blobs[0]]
inp_blobs[2] = w_grad
# add layer to the model
if kwargs:
new_layer = layer(inp_blobs, out_blobs, **kwargs)
else:
new_layer = layer(inp_blobs, out_blobs)
return new_layer