def print_net(model, namescope='gpu_0'):
"""Print the model network."""
logger.info('Printing model: {}'.format(model.net.Name()))
op_list = model.net.Proto().op
for op in op_list:
input_name = op.input
# For simplicity: only print the first output
# Not recommended if there are split layers
output_name = str(op.output[0])
op_type = op.type
op_name = op.name
if namescope is None or output_name.startswith(namescope):
# Only print the forward pass network
if output_name.find('grad') >= 0 or output_name.find('__m') >= 0:
continue
output_shape = workspace.FetchBlob(output_name).shape
first_blob = True
op_label = op_type + (op_name if op_name == '' else ':' + op_name)
suffix = ' ------- (op: {})'.format(op_label)
for j in range(len(input_name)):
if input_name[j] in model.params:
continue
input_blob = workspace.FetchBlob(input_name[j])
if isinstance(input_blob, np.ndarray):
input_shape = input_blob.shape
logger.info('{:28s}: {:20s} => {:28s}: {:20s}{}'.format(
c2_utils.UnscopeName(str(input_name[j])),
'{}'.format(input_shape),
c2_utils.UnscopeName(str(output_name)),
'{}'.format(output_shape), suffix))
if first_blob:
first_blob = False
suffix = ' ------|'
logger.info('End of model: {}'.format(model.net.Name()))
def add_stop_gradient_op_to_net(model, blob_name):
""" Add StopGradient to a defined model. The blob_name is the
name of the blob that you want to stop the gradient.
"""
if isinstance(blob_name, str):
blob_stop_grad_list = [core.ScopedBlobReference(blob_name)]
elif isinstance(blob_name, core.BlobReference):
blob_stop_grad_list = [blob_name]
elif isinstance(blob_name, list) or isinstance(blob_name, tuple):
blob_stop_grad_list = [core.ScopedBlobReference(b) for b in blob_name]
else:
raise TypeError('Argument {} must be a string or ' +
'BlobReference or list, tuple'.format(blob_name))
# Get the op's index for the stop_gradient blob
stop_grad_index = []
sorted_name = []
for i, op in enumerate(model.net._net.op):
for o in op.output:
if o in blob_stop_grad_list:
stop_grad_index.append(i)
sorted_name.append(o)
# Construct sub nets
op_index = [-1] + stop_grad_index + [len(model.net._net.op) - 1]
sub_nets = []
for i in range(len(op_index) - 1):
start_id, end_id = op_index[i] + 1, op_index[i + 1] + 1
sub_nets.append(model.net._net.op[start_id:end_id])
del model.net._net.op[:]
# Add a StopGradientOp after adding each subnet to the net
for i in range(len(sub_nets) - 1):
model.net._net.op.extend(sub_nets[i])
blob_stop_grad = c2_utils.UnscopeName(str(sorted_name[i]))
model.net.StopGradient(blob_stop_grad, blob_stop_grad)
# Add the last subnet to the net
model.net._net.op.extend(sub_nets[-1])
return model
def initialize_gpu_0_from_weights_file(model, weights_file):
"""Initialize a network with ops on GPU 0. Note that we always use GPU 0 and
rely on proper usage of CUDA_VISIBLE_DEVICES.
"""
logger.info('Loading from: {}'.format(weights_file))
ws_blobs = workspace.Blobs()
with open(weights_file, 'r') as f:
src_blobs = pickle.load(f)
if 'cfg' in src_blobs:
saved_cfg = yaml.load(src_blobs['cfg'], Loader=yamlloader.ordereddict.CLoader)
configure_bbox_reg_weights(model, saved_cfg)
if 'blobs' in src_blobs:
# Backwards compat--dictionary used to be only blobs, now they are
# stored under the 'blobs' key
src_blobs = src_blobs['blobs']
# Initialize weights on GPU 0 only
unscoped_param_names = OrderedDict() # Print these out in model order
for blob in model.params:
unscoped_param_names[c2_utils.UnscopeName(str(blob))] = True
with c2_utils.NamedCudaScope(0):
for unscoped_param_name in unscoped_param_names.keys():
if (unscoped_param_name.find(']_') >= 0 and
unscoped_param_name not in src_blobs):
# Special case for sharing initialization from a pretrained
# model:
# If a blob named '_[xyz]_foo' is in model.params and not in
# the initialization blob dictionary, then load source blob
# 'foo' into destination blob '_[xyz]_foo'
src_name = unscoped_param_name[
unscoped_param_name.find(']_') + 2:]
else:
src_name = unscoped_param_name
if src_name not in src_blobs:
logger.info('{:s} not found'.format(src_name))
continue
dst_name = core.ScopedName(unscoped_param_name)
has_momentum = src_name + '_momentum' in src_blobs
has_momentum_str = ' [+ momentum]' if has_momentum else ''
logger.info('{:s}{:} loaded from weights file into {:s}: {}'.
format(
src_name, has_momentum_str,
dst_name, src_blobs[src_name].shape))
if dst_name in ws_blobs:
# If the blob is already in the workspace, make sure that it
# matches the shape of the loaded blob
ws_blob = workspace.FetchBlob(dst_name)
assert ws_blob.shape == src_blobs[src_name].shape, \
('Workspace blob {} with shape {} does not match '
'weights file shape {}').format(
src_name,
ws_blob.shape,
src_blobs[src_name].shape)
workspace.FeedBlob(
dst_name,
src_blobs[src_name].astype(np.float32, copy=False))
if has_momentum:
workspace.FeedBlob(
dst_name + '_momentum',
src_blobs[src_name + '_momentum'].astype(
np.float32, copy=False))
# We preserve blobs that are in the weights file but not used by the current
# model. We load these into CPU memory under the '__preserve__/' namescope.
# These blobs will be stored when saving a model to a weights file. This
# feature allows for alternating optimization of Faster R-CNN in which blobs
# unused by one step can still be preserved forward and used to initialize
# another step.
for src_name in src_blobs.keys():
if (src_name not in unscoped_param_names and
not src_name.endswith('_momentum') and
src_blobs[src_name] is not None):
with c2_utils.CpuScope():
workspace.FeedBlob(
'__preserve__/{:s}'.format(src_name), src_blobs[src_name])
logger.info(
'{:s} preserved in workspace (unused)'.format(src_name))
def AddLosses(self, losses):
if not isinstance(losses, list):
losses = [losses]
# Conversion to str allows losses to include BlobReferences
losses = [c2_utils.UnscopeName(str(l)) for l in losses]
self.losses = list(set(self.losses + losses))
def unscope_name(name):
return c2_utils.UnscopeName(name)
def initialize_gpu_from_weights_file(model,
weights_file,
gpu_id=0,
train=False):
"""Initialize a network with ops on a specific GPU.
If you use CUDA_VISIBLE_DEVICES to target specific GPUs, Caffe2 will
automatically map logical GPU ids (starting from 0) to the physical GPUs
specified in CUDA_VISIBLE_DEVICES.
"""
logger.info('Loading weights from: {}'.format(weights_file))
ws_blobs = workspace.Blobs()
with open(weights_file, 'r') as f:
src_blobs = pickle.load(f)
if 'cfg' in src_blobs:
saved_cfg = yaml.load(src_blobs['cfg'])
configure_bbox_reg_weights(model, saved_cfg)
if 'blobs' in src_blobs:
# Backwards compat--dictionary used to be only blobs, now they are
# stored under the 'blobs' key
src_blobs = src_blobs['blobs']
if cfg.DISTILLATION.DISTILLATION_ON and train:
logger.info('Loading teacher weights from: {}'.format(
teacher_cfg.TRAIN.WEIGHTS))
with open(teacher_cfg.TRAIN.WEIGHTS, 'r') as f:
teacher_blobs = pickle.load(f)
if 'blobs' in teacher_blobs:
teacher_blobs = teacher_blobs['blobs']
for k, v in teacher_blobs.items():
src_blobs['teacher/{}'.format(k)] = v
# Initialize weights on GPU gpu_id only
unscoped_param_names = OrderedDict() # Print these out in model order
for blob in model.params:
unscoped_param_names[c2_utils.UnscopeName(str(blob))] = True
with c2_utils.NamedCudaScope(gpu_id):
for unscoped_param_name in unscoped_param_names.keys():
if (unscoped_param_name.find(']_') >= 0
and unscoped_param_name not in src_blobs):
# Special case for sharing initialization from a pretrained
# model:
# If a blob named '_[xyz]_foo' is in model.params and not in
# the initialization blob dictionary, then load source blob
# 'foo' into destination blob '_[xyz]_foo'
src_name = unscoped_param_name[unscoped_param_name.find(']_') +
2:]
else:
src_name = unscoped_param_name
if src_name not in src_blobs:
logger.info('{:s} not found'.format(src_name))
continue
dst_name = core.ScopedName(unscoped_param_name)
has_momentum = src_name + '_momentum' in src_blobs
has_momentum_str = ' [+ momentum]' if has_momentum else ''
logger.debug(
'{:s}{:} loaded from weights file into {:s}: {}'.format(
src_name, has_momentum_str, dst_name,
src_blobs[src_name].shape))
if dst_name in ws_blobs:
# If the blob is already in the workspace, make sure that it
# matches the shape of the loaded blob
ws_blob = workspace.FetchBlob(dst_name)
#print(src_name,dst_name)
if ws_blob.shape != src_blobs[src_name].shape:
logger.info(
"Shape missmatch: name: {} src: {}, dst: {}".format(
dst_name, ws_blob.shape,
src_blobs[src_name].shape))
continue
assert ws_blob.shape == src_blobs[src_name].shape, \
('Workspace blob {} with shape {} does not match '
'weights file shape {}').format(
src_name,
ws_blob.shape,
src_blobs[src_name].shape)
workspace.FeedBlob(
dst_name, src_blobs[src_name].astype(np.float32, copy=False))
if has_momentum:
workspace.FeedBlob(
dst_name + '_momentum',
src_blobs[src_name + '_momentum'].astype(np.float32,
copy=False))
# We preserve blobs that are in the weights file but not used by the current
# model. We load these into CPU memory under the '__preserve__/' namescope.
# These blobs will be stored when saving a model to a weights file. This
# feature allows for alternating optimization of Faster R-CNN in which blobs
# unused by one step can still be preserved forward and used to initialize
# another step.
for src_name in src_blobs.keys():
if (src_name not in unscoped_param_names
and not src_name.endswith('_momentum')
and src_blobs[src_name] is not None):
with c2_utils.CpuScope():
workspace.FeedBlob('__preserve__/{:s}'.format(src_name),
src_blobs[src_name])
logger.debug(
'{:s} preserved in workspace (unused)'.format(src_name))
def initialize_gpu_from_weights_file(model, weights_file, gpu_id=0):
"""Initialize a network with ops on a specific GPU.
If you use CUDA_VISIBLE_DEVICES to target specific GPUs, Caffe2 will
automatically map logical GPU ids (starting from 0) to the physical GPUs
specified in CUDA_VISIBLE_DEVICES.
"""
logger.info('Loading weights from: {}'.format(weights_file))
ws_blobs = workspace.Blobs()
with open(weights_file, 'r') as f:
src_blobs = pickle.load(f)
if 'cfg' in src_blobs:
saved_cfg = yaml.load(src_blobs['cfg'])
configure_bbox_reg_weights(model, saved_cfg)
if 'blobs' in src_blobs:
# Backwards compat--dictionary used to be only blobs, now they are
# stored under the 'blobs' key
src_blobs = src_blobs['blobs']
#print(src_blobs)
# Initialize weights on GPU gpu_id only
unscoped_param_names = OrderedDict() # Print these out in model order
#print(model.params)
for blob in model.params:
unscoped_param_names[c2_utils.UnscopeName(str(blob))] = True
'''
keyname = c2_utils.UnscopeName(str(blob))
if (keyname == 'retnet_cls_pred_fpn3_w' or keyname == 'retnet_cls_pred_fpn3_b'
or keyname == 'retnet_bbox_pred_fpn3_w' or keyname == 'retnet_bbox_pred_fpn3_b'):
continue
unscoped_param_names[keyname] = True'''
with c2_utils.NamedCudaScope(gpu_id):
#print(unscoped_param_names.keys())
for unscoped_param_name in unscoped_param_names.keys():
#print(unscoped_param_name)
if (unscoped_param_name.find(']_') >= 0
and unscoped_param_name not in src_blobs):
# Special case for sharing initialization from a pretrained
# model:
# If a blob named '_[xyz]_foo' is in model.params and not in
# the initialization blob dictionary, then load source blob
# 'foo' into destination blob '_[xyz]_foo'
src_name = unscoped_param_name[unscoped_param_name.find(']_') +
2:]
else:
src_name = unscoped_param_name
#print(src_name)
if src_name not in src_blobs:
logger.info('{:s} not found'.format(src_name))
print('{} not found'.format(src_name))
#continue
if src_name == 'retnet_cls_conv_n0_fpn2_w':
src_name = 'retnet_cls_conv_n0_fpn3_w'
elif src_name == 'retnet_cls_conv_n0_fpn2_b':
src_name = 'retnet_cls_conv_n0_fpn3_b'
elif src_name == 'retnet_cls_conv_n1_fpn2_w':
src_name = 'retnet_cls_conv_n1_fpn3_w'
elif src_name == 'retnet_cls_conv_n1_fpn2_b':
src_name = 'retnet_cls_conv_n1_fpn3_b'
elif src_name == 'retnet_cls_conv_n2_fpn2_w':
src_name = 'retnet_cls_conv_n2_fpn3_w'
elif src_name == 'retnet_cls_conv_n2_fpn2_b':
src_name = 'retnet_cls_conv_n2_fpn3_b'
elif src_name == 'retnet_cls_conv_n3_fpn2_w':
src_name = 'retnet_cls_conv_n3_fpn3_w'
elif src_name == 'retnet_cls_conv_n3_fpn2_b':
src_name = 'retnet_cls_conv_n3_fpn3_b'
elif src_name == 'retnet_bbox_conv_n0_fpn2_w':
src_name = 'retnet_bbox_conv_n0_fpn3_w'
elif src_name == 'retnet_bbox_conv_n0_fpn2_b':
src_name = 'retnet_bbox_conv_n0_fpn3_b'
elif src_name == 'retnet_bbox_conv_n1_fpn2_w':
src_name = 'retnet_bbox_conv_n1_fpn3_w'
elif src_name == 'retnet_bbox_conv_n1_fpn2_b':
src_name = 'retnet_bbox_conv_n1_fpn3_b'
elif src_name == 'retnet_bbox_conv_n2_fpn2_w':
src_name = 'retnet_bbox_conv_n2_fpn3_w'
elif src_name == 'retnet_bbox_conv_n2_fpn2_b':
src_name = 'retnet_bbox_conv_n2_fpn3_b'
elif src_name == 'retnet_bbox_conv_n3_fpn2_w':
src_name = 'retnet_bbox_conv_n3_fpn3_w'
elif src_name == 'retnet_bbox_conv_n3_fpn2_b':
src_name = 'retnet_bbox_conv_n3_fpn3_b'
else:
continue
dst_name = core.ScopedName(unscoped_param_name)
has_momentum = src_name + '_momentum' in src_blobs
has_momentum_str = ' [+ momentum]' if has_momentum else ''
logger.debug(
'{:s}{:} loaded from weights file into {:s}: {}'.format(
src_name, has_momentum_str, dst_name,
src_blobs[src_name].shape))
if dst_name in ws_blobs:
# If the blob is already in the workspace, make sure that it
# matches the shape of the loaded blob
ws_blob = workspace.FetchBlob(dst_name)
print(dst_name, ws_blob.shape)
print(src_name, src_blobs[src_name].shape)
assert ws_blob.shape == src_blobs[src_name].shape, \
('Workspace blob {} with shape {} does not match '
'weights file shape {}').format(
src_name,
ws_blob.shape,
src_blobs[src_name].shape)
workspace.FeedBlob(
dst_name, src_blobs[src_name].astype(np.float32, copy=False))
if has_momentum:
workspace.FeedBlob(
dst_name + '_momentum',
src_blobs[src_name + '_momentum'].astype(np.float32,
copy=False))
# We preserve blobs that are in the weights file but not used by the current
# model. We load these into CPU memory under the '__preserve__/' namescope.
# These blobs will be stored when saving a model to a weights file. This
# feature allows for alternating optimization of Faster R-CNN in which blobs
# unused by one step can still be preserved forward and used to initialize
# another step.
'''for src_name in src_blobs.keys():
