def Snapshot(
tensors,
filename,
prefix='',
suffix='.bin',
format='default',
):
"""Snapshot tensors into a binary file.
Parameters
----------
tensors : list of Tensor or Tensor
The tensors to be wrote.
filename : str
The name of this binary file.
prefix : str
The prefix of this binary file.
suffix : str
The suffix of this binary file.
format : str
The format of this binary file.
Returns
-------
None
Notes
-----
The full file path will be: ``prefix`` + ``filename`` + ``suffix``.
Available formats: ['default', 'caffe'].
"""
file_path = prefix + filename + suffix
if mpi.Is_Init():
if not mpi.AllowSnapshot(): return
file_path = file_path + '.rank.{}'.format(mpi.Rank())
dir = os.path.split(file_path)[0]
if len(dir) > 0 and not os.path.exists(dir): os.makedirs(dir)
if format == 'default':
state_dict = {}
for tensor in tensors:
state_dict[tensor.name] = FetchTensor(tensor)
with open(file_path, 'wb') as f:
pickle.dump(state_dict, f, pickle.HIGHEST_PROTOCOL)
logging.info('Snapshot Model@: ' + file_path)
logging.info('Model Format: Pickle')
elif format is 'caffe':
names = [tensor.name for tensor in tensors]
_C.Snapshot(file_path, names, 1)
else:
raise TypeError('Unknown binary format: {}'.format(format))
def Snapshot(
tensors,
filename,
prefix='',
suffix='.bin',
format='pickle',
):
"""Serialize tensors into a binary file.
The filename is formatted as:
``prefix`` + ``filename`` + ``suffix``
Parameters
----------
tensors : list of Tensor or Tensor
The tensors to be wrote.
filename : str
The name of this binary file.
prefix : str, optional, default=''
The prefix of this binary file.
suffix : str, optional, default='.bin'
The suffix of this binary file.
format : {'pickle', 'caffe'}, optional
The format of this binary file.
Returns
-------
None
"""
file_path = prefix + filename + suffix
if _mpi.Is_Init():
if not _mpi.AllowSnapshot(): return
file_path = file_path + '.rank.{}'.format(_mpi.Rank())
dir = os.path.split(file_path)[0]
if len(dir) > 0 and not os.path.exists(dir): os.makedirs(dir)
if format == 'pickle':
state_dict = {}
for tensor in tensors:
state_dict[tensor.name] = FetchTensor(tensor)
with open(file_path, 'wb') as f:
pickle.dump(state_dict, f, pickle.HIGHEST_PROTOCOL)
_logging.info('Snapshot Model@: ' + file_path)
_logging.info('Model Format: Pickle')
elif format == 'caffe':
names = [tensor.name for tensor in tensors]
get_default_workspace().Snapshot(file_path, names, 1)
else:
raise TypeError('Unknown binary format: ' + format)
def snapshot(self):
if mpi.Is_Init():
if not mpi.AllowSnapshot(): return
net = self.solver.net
infix = ('_' + cfg.TRAIN.SNAPSHOT_INFIX
if cfg.TRAIN.SNAPSHOT_INFIX != '' else '')
filename = (self.solver_param.snapshot_prefix + infix +
'_iter_{:d}'.format(self.solver.iter) + '.caffemodel')
filename = os.path.join(self.output_dir, filename)
net.save(str(filename))
logger.info('Wrote snapshot to: {:s}'.format(filename))
return filename