def add_update_defs(graph_def, optimizer):
"""Add the update defs."""
grads, update_defs = [], []
extra_arguments = optimizer._extra_kwargs
extra_arguments['handle'] = optimizer._op_handle
# Generate op defs according to the collected updates
current_ws = workspace.get_workspace()
for (param, grad), arguments in optimizer._param_group:
if current_ws.has_tensor(grad):
grads.append(grad)
arguments = dict(arguments, **extra_arguments)
update_defs.append(
proto_util.make_operator_def(op_type=optimizer._op_type,
inputs=[grad],
outputs=[param],
name=OpDef.get_name(),
**arguments))
else:
logging.info('Skip to update Tensor({}).'.format(param))
# Insert a reduce def if the process group is found.
process_group = optimizer._process_group
if process_group is not None:
update_defs.insert(
0,
proto_util.make_operator_def(op_type='Collective',
inputs=grads,
outputs=grads,
name=OpDef.get_name(),
operation='MEAN',
communication='ALLREDUCE',
**process_group.arguments))
graph_def.op.extend(update_defs)
def load_weights_from_pickle(f, layer, verbose=False):
ws = workspace.get_workspace()
weight_dict = six.moves.pickle.load(f)
for weight in layer.weights:
name = weight.name
if name in weight_dict:
value = weight_dict[name]
value_shape = list(value.shape)
weight_shape = list(weight.shape)
if value_shape != weight_shape:
raise ValueError(
'Shape of weight({}) is ({}), \n'
'While load from shape of ({}).'
.format(name, ', '.join(
[str(d) for d in weight_shape]),
', '.join([str(d) for d in value_shape]))
)
weight_impl = ws.GetTensor(weight.id)
if weight_impl is not None:
weight_impl.FromNumpy(value.copy())
if verbose:
logging.info(
'Weight({}) loaded, Size: ({})'
.format(name, ', '.join([str(d) for d in value_shape])))
else:
logging.warning(
'Weight({}) is not created '
'in current workspace. Skip.'.format(name))
def create_graph(self, graph_def):
"""Create a graph."""
cfg = config.config()
if cfg.graph_verbosity == 2:
msg = '\n' + str(graph_def)[:-1]
logging.info('\ngraph {' + msg.replace('\n', '\n ') + '\n}\n')
return self._impl.CreateGraph(serialization.serialize_proto(graph_def),
cfg.graph_verbosity == 1)
def cleanup():
def terminate(processes):
for p in processes:
p.terminate()
p.join()
terminate(self._workers)
if rank == 0:
logging.info('Terminate DataWorker.')
terminate(self._readers)
if rank == 0:
logging.info('Terminate DataReader.')
def _get_learning_rate(self):
"""Get the learning rate based on preset policy."""
policy = self._proto.lr_policy
if policy == "step":
new_step = int(self.iter / self._proto.stepsize)
if self._current_step != new_step:
new_lr = self._proto.base_lr * pow(self._proto.gamma, new_step)
self._current_step = new_step
self.base_lr = new_lr
elif policy == 'multistep':
if self._current_step < len(self._proto.stepvalue) \
and self.iter >= self._proto.stepvalue[self._current_step]:
self._current_step = self._current_step + 1
logging.info(
'MultiStep Status: Iteration {}, step = {}'.format(
self.iter, self._current_step))
new_lr = (self._proto.base_lr *
pow(self._proto.gamma, self._current_step))
self.base_lr = new_lr
elif policy == 'multifixed':
stage_lrs = self._proto.stage_lr
stage_iters = self._proto.stage_iter
if self.iter < stage_iters[self._current_step]:
self.base_lr = stage_lrs[self._current_step]
else:
if self._current_step + 1 < len(stage_iters):
self._current_step = self._current_step + 1
logging.info(
'MultiFixed Status: Iteration {}, stage = {}'.format(
self.iter, self._current_step))
self.base_lr = stage_lrs[self._current_step]
elif policy == 'inv':
power = self._proto.power
gamma = self._proto.gamma
self.base_lr = (self._proto.base_lr *
pow(1. + gamma * self.iter, -power))
elif policy == 'poly':
power = self._proto.power
max_iter = self._proto.max_iter
self.base_lr = (self._proto.base_lr *
pow(1. - float(self.iter) / max_iter, power))
def create_graph(self, graph_def):
"""Create the graph.
Parameters
----------
graph_def : GraphDef
The ``GraphDef`` protocol buffer.
Returns
-------
str
The graph name.
"""
cfg = config.config()
if cfg.graph_verbosity == 2:
msg = '\n' + str(graph_def)[:-1]
logging.info('\ngraph {' + msg.replace('\n', '\n ') + '\n}\n')
return self.CreateGraph(
serialization.serialize_proto(graph_def),
cfg.graph_verbosity == 1)
def test(self, test_idx):
"""Test the specific net.
Parameters
----------
test_idx : int
The index of test net.
"""
net = self._test_nets[test_idx]
net.copy_from(self._net.to_proto())
test_iter = self._proto.test_iter[test_idx]
test_scores = collections.defaultdict(float)
for iter in range(test_iter):
net.forward()
for key in net.outputs:
test_scores[key] += float(net.blobs[key].data.numpy().mean())
logging.info('Iteration {}, Test net #{}'.format(self.iter, test_idx))
for i, (key, score) in enumerate(test_scores.items()):
logging.info(' ' * 4 + 'Test net output #%d(%s): %.4f' %
(i, key, score / test_iter))
def step(self, num_iterations=1):
"""Step the train net.
Parameters
----------
num_iterations : int, optional, default=1
The number of iterations to step.
"""
end_iter = self.iter + num_iterations
tic = time.time()
while self.iter < end_iter:
# Test if necessary.
if (self._proto.test_interval > 0
and self.iter % self._proto.test_interval == 0):
if (self.iter == 0
and self._proto.test_initialization) or self.iter != 0:
for test_idx in range(len(self._test_nets)):
self.test(test_idx)
# Forward and backward pass.
self._net.forward()
# Display iteration info.
if self._is_root and self._proto.display:
if self.iter % self._proto.display == 0:
logging.info('Iteration %d, lr = %f, time = %.2fs' %
(self.iter, self.base_lr, time.time() - tic))
tic = time.time()
for i, key in enumerate(self.net.outputs):
value = self.net.blobs[key].data.numpy().mean()
logging.info(' ' * 4 +
'Train net output #{}({}): {}'.format(
i, key, value))
# Apply Update.
self._get_learning_rate()
self._apply_update()
self.iter = self.iter + 1
# Snapshot if necessary.
if self._proto.snapshot:
if self.iter % self._proto.snapshot == 0:
self.snapshot()
def export_to(self, name=None, export_dir='./'):
"""Export the graph into a text file.
Parameters
----------
name : str
The optional graph name.
export_dir : str
The directory to export the file.
"""
if not os.path.exists(export_dir):
try:
os.makedirs(export_dir)
except Exception:
raise ValueError('The given directory can not be created.')
graph_def = copy.deepcopy(self.graph_def)
graph_def.name = self.graph_def.name if name is None else name
path = os.path.join(export_dir, graph_def.name + '.graph')
with open(path, 'w') as f:
f.write(str(graph_def))
logging.info('Export meta graph into: {}'.format(path))
def load_weights(self, filepath, format=None, skip=False, verbose=False):
"""Load model weights from a binary file.
Parameters
----------
filepath : str
The path of weights file.
format : {'hdf5', 'npz', 'pkl', 'npz_dict'}, optional
The optional saving format.
skip : bool, optional, default=False
``True`` to skip the modules which is not found.
verbose: bool, optional, default=False
``True`` to print the matched weights.
"""
if not os.path.exists(filepath):
raise FileNotFoundError("File {} doesn't exist.".format(filepath))
if format is None:
format = filepath.split('.')[-1]
if format == 'hdf5' or format == 'h5':
matched_info = files.load_hdf5_to_weights(filepath, self, skip)
elif format == 'npz':
matched_info = files.load_and_assign_npz(filepath, self)
elif format == 'pkl':
matched_info = files.load_and_assign_pkl_dict(filepath, self, skip)
elif format == 'npz_dict':
matched_info = files.load_and_assign_npz_dict(filepath, self, skip)
else:
raise ValueError(
"Saving format should be in ('hdf5', 'npz', 'pkl', 'npz_dict').\n"
"Format <%s> is not supported." % format)
if verbose:
for info in matched_info:
logging.info('Weight({}) loaded, Size: ({})'.format(
info[0], ', '.join([str(d) for d in info[1]])))
def from_proto(self, proto):
"""Deserialize from the proto.
Parameters
----------
proto : LayerParameter
The ``LayerParameter`` protocol buffer.
"""
for i in range(len(self._blobs)):
if i < len(proto.blobs):
blob_proto = proto.blobs[i]
if len(blob_proto.data) > 0:
value = numpy.array(blob_proto.data, dtype='float32')
elif len(blob_proto.double_data) > 0:
value = numpy.array(blob_proto.double_data, dtype='float64')
else:
raise ValueError('Neither <data> or <double_data> in blob proto.')
if len(blob_proto.shape.dim) > 0:
value = value.reshape([dim for dim in blob_proto.shape.dim])
self._blobs[i]['data'].set_value(value)
logging.info('Blob({}/param:{}) loaded, shape: {}, size: {}'
.format(self._name, i, value.shape, value.size))
def test(self, test_idx):
"""Test the specific net.
Parameters
----------
test_idx : int
The idx of test net.
"""
test_score, output_id = [], []
net = self._test_nets[test_idx]
test_iter = self._param.test_iter[test_idx]
for iter in range(test_iter):
net.forward()
if not self._is_root:
continue
if iter == 0:
for key in net.outputs:
values = net.blobs[key].data.get_value().flatten()
for idx, value in enumerate(values):
test_score.append(value)
output_id.append(key)
else:
i = 0
for key in net.outputs:
values = net.blobs[key].data.get_value().flatten()
for idx, value in enumerate(values):
test_score[i] += value
i += 1
logging.info('Iteration {}, Test net #{}'.format(self.iter, test_idx))
for i, score in enumerate(test_score):
logging.info(
' ' * 10 + 'Test net output #%d(%s): %.4f'
% (i, output_id[i], score / test_iter))
def __init__(self, cuda_engine, device_id=0):
"""Create an ``Engine``.
Parameters
----------
cuda_engine : tensorrt.ICudaEngine
The built cuda engine.
device_id : int, optional, default=0
The index of executing device.
"""
# Create executing resources.
self._cuda_engine = cuda_engine
self._device_id = device_id
self._context = cuda_engine.create_execution_context()
self._stream = driver.Stream(0)
# Create bindings.
num_binding = self._cuda_engine.num_bindings
self._bindings = [
Binding(cuda_engine, self._context, i, device_id)
for i in range(num_binding)
]
self._inputs = [b for b in self._bindings if b.is_input]
self._outputs = [b for b in self._bindings if not b.is_input]
# Report the engine info.
logging.info('TensorRT engine built.')
binding_info = 'InputInfo: {\n'
for b in self._inputs:
binding_info += ' * Binding("{}", shape={}, dtype={})\n' \
.format(b.name, b.shape, b.dtype)
logging.info(binding_info + '}')
binding_info = 'OutputInfo: {\n'
for b in self._outputs:
binding_info += ' * Binding("{}", shape={}, dtype={})\n' \
.format(b.name, b.shape, b.dtype)
logging.info(binding_info + '}')
def step(self, num_iterations=1):
"""Step the train net.
Parameters
----------
num_iterations : int, optional, default=1
The number of iterations to step.
"""
start_step = self.iter
stop_step = start_step + num_iterations
loss_vec, smoothed_loss = [], 0.
tic = time.time()
while self.iter < stop_step:
# Test if necessary.
if self._is_root and self._param.test_interval > 0 and \
self.iter % self._param.test_interval == 0:
if (self.iter == 0 and self._param.test_initialization) or \
self.iter != 0:
for test_idx in range(len(self._test_nets)):
self.test(test_idx)
# Forward, backward and compute loss.
loss = 0.
for i in range(self._param.iter_size):
self._net.forward_backward()
if self._is_root:
for e in self.net.losses:
values = e.get_value().flatten()
for v in values:
loss += v
if self._is_root:
loss /= self._param.iter_size
if len(loss_vec) < self._param.average_loss:
loss_vec.append(loss)
smoothed_loss *= (len(loss_vec) - 1)
smoothed_loss += loss
smoothed_loss /= len(loss_vec)
else:
idx = (self.iter - start_step) % self._param.average_loss
smoothed_loss += ((loss - loss_vec[idx]) / self._param.average_loss)
loss_vec[idx] = loss
# Apply Update.
self._get_learning_rate()
self._apply_update()
# Display iteration info.
if self._is_root and self._param.display:
if self.iter % self._param.display == 0:
logging.info(
'Iteration %d, lr = %s, loss = %f, time = %.2fs'
% (self.iter, str(self.base_lr), smoothed_loss, time.time() - tic))
tic = time.time()
for idx, net_output in enumerate(self.net.outputs):
values = self.net.blobs[net_output].data.get_value().flatten()
for v in values:
logging.info(
' ' * 10 + 'Train net output #{}({}): {}'
.format(idx, net_output, v))
self.iter = self.iter + 1
# Snapshot if necessary.
if self._param.snapshot:
if self.iter % self._param.snapshot == 0:
self.snapshot()