def apply_optimizers(
self,
train_net,
train_init_net,
grad_map,
blob_to_device=None,
):
CPU = muji.OnCPU()
# if given, blob_to_device is a map from blob to device_option
blob_to_device = blob_to_device or {}
for param, optimizer in viewitems(self.param_to_optim):
assert optimizer is not None, \
"default optimizer must have been set in add_layer"
# note that not all params has gradient and thus we sent None if
# gradient does not exists
device = get_param_device(
param,
grad_map.get(str(param)),
param_to_device=blob_to_device,
default_device=CPU,
)
if device is not None:
# extra info is not applicable for optimizers
del device.extra_info[:]
with core.DeviceScope(device):
optimizer(
train_net, train_init_net, param, grad_map.get(str(param)))
def apply_regularizers_after_optimizer(
self,
train_net,
train_init_net,
grad_map,
blob_to_device=None,
):
logger.info("apply regularizer after optimizer")
CPU = muji.OnCPU()
# if given, blob_to_device is a map from blob to device_option
blob_to_device = blob_to_device or {}
for param, regularizer in viewitems(self.param_to_reg):
if regularizer is None:
continue
assert isinstance(regularizer, Regularizer)
logger.info("add regularizer {0} for param {1} to optimizer".format(regularizer, param))
device = get_param_device(
param,
grad_map.get(str(param)),
param_to_device=blob_to_device,
default_device=CPU,
)
with core.DeviceScope(device):
regularizer(
train_net, train_init_net, param, grad=grad_map.get(str(param)),
by=RegularizationBy.AFTER_OPTIMIZER
)
def modify_net(self,
net,
init_net=None,
grad_map=None,
blob_to_device=None,
modify_output_record=False):
p = self._p
compute_averaged_norm = self._compute_averaged_norm
row_index = self.row_index
CPU = muji.OnCPU()
# if given, blob_to_device is a map from blob to device_option
blob_to_device = blob_to_device or {}
for blob_name in self._blobs:
blob = core.BlobReference(blob_name)
assert net.BlobIsDefined(
blob
), 'blob {} is not defined in net {} whose proto is {}'.format(
blob, net.Name(), net.Proto())
if blob in blob_to_device:
device = blob_to_device[blob]
else:
device = CPU
with core.DeviceScope(device):
if row_index and row_index >= 0:
blob = net.Slice(
[blob],
net.NextScopedBlob(prefix=blob +
'_row_{0}'.format(row_index)),
starts=[row_index, 0],
ends=[row_index + 1, -1])
cast_blob = net.Cast(blob,
net.NextScopedBlob(prefix=blob +
'_float'),
to=core.DataType.FLOAT)
norm_name = net.NextScopedBlob(prefix=blob +
self._field_name_suffix)
norm = net.LpNorm(cast_blob,
norm_name,
p=p,
average=compute_averaged_norm)
if self._logging_frequency >= 1:
net.Print(norm, [], every_n=self._logging_frequency)
if modify_output_record:
output_field_name = str(blob) + self._field_name_suffix
output_scalar = schema.Scalar((np.float, (1, )), norm)
if net.output_record() is None:
net.set_output_record(
schema.Struct((output_field_name, output_scalar)))
else:
net.AppendOutputRecordField(output_field_name,
output_scalar)
def apply_regularizers_after_optimizer(
self,
train_net,
train_init_net,
grad_map,
blob_to_device=None,
):
CPU = muji.OnCPU()
# if given, blob_to_device is a map from blob to device_option
blob_to_device = blob_to_device or {}
for param, regularizer in viewitems(self.param_to_reg):
if regularizer is None or not regularizer.apply_after_optimizer:
continue
assert isinstance(regularizer, Regularizer)
device = get_param_device(
param,
grad_map.get(str(param)),
param_to_device=blob_to_device,
default_device=CPU,
)
with core.DeviceScope(device):
regularizer(
train_net, train_init_net, param, grad_map.get(str(param)))
