def __init__(self,
backbone,
twin,
twin_load_from,
neck=None,
bbox_head=None,
pair_module=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(TwinV2SingleStageDetector, self).__init__()
self.backbone = builder.build_backbone(backbone)
if neck is not None:
self.neck = builder.build_neck(neck)
self.neck_first = True
if pair_module is not None:
self.pair_module = builder.build_pair_module(
pair_module)
if hasattr(self.pair_module, 'neck_first'):
self.neck_first = self.pair_module.neck_first
self.bbox_head = builder.build_head(bbox_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
# Build twin model
print(" Loading Twin's weights...")
self.twin = build_detector(twin, train_cfg=self.train_cfg, test_cfg=self.test_cfg)
# self.twin = MMDataParallel(twin, device_ids=[0]).cuda() # TODO check device id?
load_checkpoint(self.twin, twin_load_from, map_location='cpu', strict=False, logger=None)
print(" Finished loading twin's weight. ")
self.twin.eval()
# memory cache for testing
self.key_feat = None
def __init__(self,
det_load_from,
backbone,
neck=None,
bbox_head=None,
tx_head=None,
temporal_module=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(SeqTxSingleStage, self).__init__()
self.backbone = builder.build_backbone(backbone)
if neck is not None:
self.neck = builder.build_neck(neck)
assert temporal_module is None
self.bbox_head = builder.build_head(bbox_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
# self.init_weights(pretrained=pretrained)
load_checkpoint(self,
det_load_from,
map_location='cpu',
strict=False,
logger=None)
self.backbone.eval()
self.neck.eval()
self.bbox_head.eval()
self.tx_head = builder.build_head(tx_head)
self.test_seq_len = 4
self.memory = []
def test_load_classes_name():
import os
import tempfile
from mmcv.runner import load_checkpoint, save_checkpoint
checkpoint_path = os.path.join(tempfile.gettempdir(), 'checkpoint.pth')
model = Model()
save_checkpoint(model, checkpoint_path)
checkpoint = load_checkpoint(model, checkpoint_path)
assert 'meta' in checkpoint and 'CLASSES' not in checkpoint['meta']
model.CLASSES = ('class1', 'class2')
save_checkpoint(model, checkpoint_path)
checkpoint = load_checkpoint(model, checkpoint_path)
assert 'meta' in checkpoint and 'CLASSES' in checkpoint['meta']
assert checkpoint['meta']['CLASSES'] == ('class1', 'class2')
model = Model()
wrapped_model = DDPWrapper(model)
save_checkpoint(wrapped_model, checkpoint_path)
checkpoint = load_checkpoint(wrapped_model, checkpoint_path)
assert 'meta' in checkpoint and 'CLASSES' not in checkpoint['meta']
wrapped_model.module.CLASSES = ('class1', 'class2')
save_checkpoint(wrapped_model, checkpoint_path)
checkpoint = load_checkpoint(wrapped_model, checkpoint_path)
assert 'meta' in checkpoint and 'CLASSES' in checkpoint['meta']
assert checkpoint['meta']['CLASSES'] == ('class1', 'class2')
# remove the temp file
os.remove(checkpoint_path)
def test_load_checkpoint():
import os
import tempfile
import re
class PrefixModel(nn.Module):
def __init__(self):
super().__init__()
self.backbone = Model()
pmodel = PrefixModel()
model = Model()
checkpoint_path = os.path.join(tempfile.gettempdir(), 'checkpoint.pth')
# add prefix
torch.save(model.state_dict(), checkpoint_path)
state_dict = load_checkpoint(pmodel,
checkpoint_path,
revise_keys=[(r'^', 'backbone.')])
for key in pmodel.backbone.state_dict().keys():
assert torch.equal(pmodel.backbone.state_dict()[key], state_dict[key])
# strip prefix
torch.save(pmodel.state_dict(), checkpoint_path)
state_dict = load_checkpoint(model,
checkpoint_path,
revise_keys=[(r'^backbone\.', '')])
for key in state_dict.keys():
key_stripped = re.sub(r'^backbone\.', '', key)
assert torch.equal(model.state_dict()[key_stripped], state_dict[key])
os.remove(checkpoint_path)
def worker_func(model_cls, model_kwargs, checkpoint, dataset, data_func,
gpu_id, idx_queue, result_queue):
model = model_cls(**model_kwargs)
load_checkpoint(model, checkpoint, map_location='cpu')
torch.cuda.set_device(gpu_id)
model.cuda()
model.eval()
with torch.no_grad():
while True:
idx = idx_queue.get()
data = dataset[idx]
result = model(**data_func(data, gpu_id))
result_queue.put((idx, result))
def after_build_model(self, model):
"""Remove all pruned channels in finetune stage.
We add this function to ensure that this happens before DDP's
optimizer's initialization
"""
if not self.pruning:
for name, module in model.named_modules():
add_pruning_attrs(module)
assert self.deploy_from, 'You have to give a ckpt' \
'containing the structure information of the pruning model'
load_checkpoint(model, self.deploy_from)
deploy_pruning(model)
def extract_feats(model,
datasets,
cfg,
save_dir,
data_split='train',
logger=None):
load_checkpoint(model, cfg.load_from, 'cpu', False, logger)
fg_th = cfg.train_cfg.rcnn.assigner.pos_iou_thr
bg_th = cfg.train_cfg.rcnn.assigner.neg_iou_thr
logger.info('load checkpoint from %s', cfg.load_from)
logger.info(
f'fg_iou_thr {fg_th} bg_iou_thr {bg_th} data_split {data_split} save_dir {save_dir}'
)
model.eval()
model = model.cuda()
data_loaders = [
build_dataloader(ds,
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
1,
dist=False) for ds in datasets
]
feats = []
labels = []
for index, data in enumerate(data_loaders[0]):
bbox_feats, bbox_labels, bboxes = model.feats_extract(
data['img'].data[0], data['img_meta'].data[0],
data['gt_bboxes'].data[0], data['gt_labels'].data[0])
logger.info(
f"{index:05}/{len(data_loaders[0])} feats shape - {bbox_feats.shape}"
)
feats.append(bbox_feats.data.cpu().numpy())
labels.append(bbox_labels.data.cpu().numpy())
del data, bbox_feats, bbox_labels, bboxes
feats = np.concatenate(feats)
labels = np.concatenate(labels)
split = f'{fg_th}_{bg_th}'
np.save(f'{save_dir}/{data_split}_{split}_feats.npy', feats)
np.save(f'{save_dir}/{data_split}_{split}_labels.npy', labels)
# import pdb; pdb.set_trace()
print(f"{labels.shape} num of features")
def __init__(self,
backbone,
midrange,
midrange_load_from,
middle='B', # 'S'
neck=None,
bbox_head=None,
pair_module=None,
pair_module2=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(LMPSingleStageDetector, self).__init__()
self.backbone = builder.build_backbone(backbone)
if neck is not None:
self.neck = builder.build_neck(neck)
if pair_module is not None:
self.pair_module = builder.build_pair_module(
pair_module)
if pair_module2 is not None:
self.pair_module2 = builder.build_pair_module(
pair_module2)
self.bbox_head = builder.build_head(bbox_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
# Build MidRange model
print(" Loading MidRange's weights...")
self.midrange = build_detector(
midrange, train_cfg=self.train_cfg, test_cfg=self.test_cfg)
load_checkpoint(self.midrange, midrange_load_from, map_location='cpu',
strict=False, logger=None)
self.midrange.eval()
self.middle = middle
# memory cache for testing
self.test_interval = test_cfg.get('test_interval', 10)
self.memory_size = test_cfg.get('memory_size', 1)
self.key_feat_pre = None # This would not be used in simple_test_b
self.key_feat_post = None
# I'm sorry I have to be explicit
self.g_ref_list_list = None
self.phi_ref_list_list = None
def __init__(self,
runner,
filename=None,
decay=0.9998,
out_dir=None,
interval=-1,
save_optimizer=True,
max_keep_ckpts=-1,
meta=None,
device='',
**kwargs):
# setting checkpoints
self.interval = interval
self.out_dir = out_dir
self.save_optimizer = save_optimizer
self.max_keep_ckpts = max_keep_ckpts
self.args = kwargs
self.create_symlink = True
self.filename = filename
self.meta = meta
# make a copy of the model for accumulating moving average of weights
self.ema = deepcopy(runner.model)
if self.filename:
runner.logger.info('load EMA checkpoint for EMA model from %s',
filename)
load_checkpoint(self.ema,
self.filename,
map_location='cpu',
strict=False)
self.ema.eval()
# self.updates = 0 # number of EMA updates
self.updates = runner.iter # number of EMA updates
self.decay = lambda x: decay * (1 - math.exp(
-x / 2000)) # decay exponential ramp (to help early epochs)
self.device = device # perform ema on different device from model if set
if device:
self.ema.to(device=device)
for p in self.ema.parameters():
p.requires_grad_(False)
runner.ema = self
def load_checkpoint(self, filename, map_location='cpu', strict=True):
self.logger.info('load checkpoint from %s', filename)
if filename.startswith(('http://', 'https://')):
url = filename
filename = '../' + url.split('/')[-1]
if get_dist_info()[0] == 0:
if osp.isfile(filename):
os.system('rm ' + filename)
os.system('wget -N -q -P ../ ' + url)
dist.barrier()
return load_checkpoint(self.model, filename, map_location, strict,
self.logger)
def prune_mask_rcnn_only(args: PruneParams):
"""
Just prune without retraining.
Args:
args: (PruneParams).
Returns: (MaskRCNN) pruned model in cuda.
"""
cfg = mmcv.Config.fromfile(args.config)
model = build_detector(cfg.model,
train_cfg=cfg.get('train_cfg'),
test_cfg=cfg.get('test_cfg'))
assert cfg.model['type'] == 'MaskRCNN', 'model type should be MaskRCNN!'
# load checkpoint
checkpoint = cp.load_checkpoint(model=model, filename=args.checkpoint)
num_before = sum([p.nelement() for p in model.backbone.parameters()])
print('Before pruning, Backbone Params = %.2fM' % (num_before / 1E6))
# PRUNE FILTERS
# func = {"ResNet50": prune_resnet50, "ResNet101": prune_resnet101}
assert args.backbone in ['ResNet50', 'ResNet101'], "Wrong backbone type!"
skip = {
'ResNet34': [2, 8, 14, 16, 26, 28, 30, 32],
'ResNet50': [2, 11, 20, 23, 89, 92, 95, 98],
'ResNet101': [2, 11, 20, 23, 89, 92, 95, 98]
}
pf_cfg, new_backbone = prune_top2_layers(arch=args.backbone,
net=model.backbone,
skip_block=skip[args.backbone],
prs=str2list(args.prs),
cuda=True)
model.backbone = new_backbone
num_after = sum([p.nelement() for p in model.backbone.parameters()])
print('After pruning: Backbone Params = %.2fM' % (num_after / 1E6))
print("Prune rate: %.2f%%" % ((num_before - num_after) / num_before * 100))
# replace checkpoint['state_dict']
checkpoint['state_dict'] = cp.weights_to_cpu(cp.get_state_dict(model))
mmcv.mkdir_or_exist(osp.dirname(args.result_path))
# save and immediately flush buffer
torch.save(checkpoint, args.result_path)
with open(args.result_path.split('.')[0] + '_cfg.txt', 'w') as f:
f.write(str(pf_cfg))
def worker_func_art(model_cls, model_kwargs, checkpoint, dataset, data_func,
gpu_id, idx_queue, result_queue, post_processor,
img_prefix, show=True, show_path=None):
""" store the img_name, img_shape, ori_shape in data_queue
return single_pred_results.
"""
model = model_cls(**model_kwargs)
load_checkpoint(model, checkpoint, map_location='cpu')
torch.cuda.set_device(gpu_id)
model.cuda()
model.eval()
with torch.no_grad():
while True:
idx = idx_queue.get()
data = dataset[idx]
data_dict = data_func(data, gpu_id)
img_metas = data_dict['img_meta'][0]
img_metas_0 = img_metas[0]
is_aug = bool(len(data['img']) > 1)
h, w, _ = img_metas_0['img_shape']
ori_h, ori_w, _ = img_metas_0['ori_shape']
filename = img_metas_0['filename']
img_name = osp.splitext(filename)[0].replace('gt_', 'res_')
scales = (ori_w * 1.0 / w, ori_h * 1.0 / h)
result = model(**data_func(data, gpu_id))
if isinstance(result, tuple):
bbox_result, segm_result = result
else:
bbox_result, segm_result = result, None
vs_bbox_result = np.vstack(bbox_result)
if segm_result is None:
pred_bboxes, pred_bbox_scores = [], []
else:
segm_scores = np.asarray(vs_bbox_result[:, -1])
segms = mmcv.concat_list(segm_result)
# the bboxes returned by processor are fit to the original images.
if not is_aug:
# print('1-metas:{:d}'.format(len(img_metas)))
# single simple test the predicted mask use the size of rescaled img.
pred_bboxes, pred_bbox_scores = post_processor.process(segms, segm_scores,
mask_shape=img_metas_0['img_shape'],
scale_factor=scales)
else:
# aug test the predicted mask use the size of original img
pred_bboxes, pred_bbox_scores = post_processor.process(segms, segm_scores,
mask_shape=img_metas_0['ori_shape'],
scale_factor=(1.0, 1.0))
# save the results.
single_pred_results = []
for pred_bbox, pred_bbox_score in zip(pred_bboxes, pred_bbox_scores):
pred_bbox = np.asarray(pred_bbox).reshape((-1, 2)).astype(np.int32)
pred_bbox = pred_bbox.tolist()
single_bbox_dict = {
"points": pred_bbox,
"confidence": float(pred_bbox_score),
}
single_pred_results.append(single_bbox_dict)
pred_result = {
"img_name": img_name,
"single_pred_results": single_pred_results
}
result_queue.put((idx, pred_result))
if show:
img = cv2.imread(osp.join(img_prefix, filename))
for idx in range(len(single_pred_results)):
bbox = np.asarray(single_pred_results[idx]["points"]).reshape(-1, 2).astype(np.int64)
cv2.drawContours(img, [bbox], -1, (0, 255, 0), 2)
cv2.imwrite(osp.join(show_path, filename), img)
def test_load_checkpoint_metadata():
import os
import tempfile
from mmcv.runner import load_checkpoint, save_checkpoint
class ModelV1(nn.Module):
def __init__(self):
super().__init__()
self.block = Block()
self.conv1 = nn.Conv2d(3, 3, 1)
self.conv2 = nn.Conv2d(3, 3, 1)
nn.init.normal_(self.conv1.weight)
nn.init.normal_(self.conv2.weight)
class ModelV2(nn.Module):
_version = 2
def __init__(self):
super().__init__()
self.block = Block()
self.conv0 = nn.Conv2d(3, 3, 1)
self.conv1 = nn.Conv2d(3, 3, 1)
nn.init.normal_(self.conv0.weight)
nn.init.normal_(self.conv1.weight)
def _load_from_state_dict(self, state_dict, prefix, local_metadata,
*args, **kwargs):
"""load checkpoints."""
# Names of some parameters in has been changed.
version = local_metadata.get('version', None)
if version is None or version < 2:
state_dict_keys = list(state_dict.keys())
convert_map = {'conv1': 'conv0', 'conv2': 'conv1'}
for k in state_dict_keys:
for ori_str, new_str in convert_map.items():
if k.startswith(prefix + ori_str):
new_key = k.replace(ori_str, new_str)
state_dict[new_key] = state_dict[k]
del state_dict[k]
super()._load_from_state_dict(state_dict, prefix, local_metadata,
*args, **kwargs)
model_v1 = ModelV1()
model_v1_conv0_weight = model_v1.conv1.weight.detach()
model_v1_conv1_weight = model_v1.conv2.weight.detach()
model_v2 = ModelV2()
model_v2_conv0_weight = model_v2.conv0.weight.detach()
model_v2_conv1_weight = model_v2.conv1.weight.detach()
ckpt_v1_path = os.path.join(tempfile.gettempdir(), 'checkpoint_v1.pth')
ckpt_v2_path = os.path.join(tempfile.gettempdir(), 'checkpoint_v2.pth')
# Save checkpoint
save_checkpoint(model_v1, ckpt_v1_path)
save_checkpoint(model_v2, ckpt_v2_path)
# test load v1 model
load_checkpoint(model_v2, ckpt_v1_path)
assert torch.allclose(model_v2.conv0.weight, model_v1_conv0_weight)
assert torch.allclose(model_v2.conv1.weight, model_v1_conv1_weight)
# test load v2 model
load_checkpoint(model_v2, ckpt_v2_path)
assert torch.allclose(model_v2.conv0.weight, model_v2_conv0_weight)
assert torch.allclose(model_v2.conv1.weight, model_v2_conv1_weight)
def load_adv_checkpoint(self, filename, map_location='cpu', strict=False):
self.logger.info('load checkpoint from %s', filename)
return load_checkpoint(self.adv_model, filename, map_location, strict,
self.logger)
def load_checkpoint(self, filename, map_location="cpu", strict=False):
self.logger.info("load checkpoint from %s", filename)
return load_checkpoint(self.model, filename, map_location, strict,
self.logger)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# update configs according to CLI args
if args.dir is not None:
if args.dir.startswith('//'):
cfg.work_dir = args.dir[2:]
else:
localhost = get_localhost().split('.')[0]
# results from server saved to /private
if 'gpu' in localhost:
output_dir = '/private/huangchenxi/mmdet/outputs'
else:
output_dir = 'work_dirs'
if args.dir.endswith('-c'):
args.dir = args.dir[:-2]
args.resume_from = search_and_delete(os.path.join(
output_dir, args.dir),
prefix=cfg.work_dir,
suffix=localhost)
cfg.work_dir += time.strftime("_%m%d_%H%M") + '_' + localhost
cfg.work_dir = os.path.join(output_dir, args.dir, cfg.work_dir)
if args.workers_per_gpu != -1:
cfg.data['workers_per_gpu'] = args.workers_per_gpu
if args.resume_from is not None:
cfg.resume_from = args.resume_from
cfg.gpus = args.gpus
if args.profiler or args.speed:
cfg.data.imgs_per_gpu = 1
if cfg.resume_from or cfg.load_from:
cfg.model['pretrained'] = None
if args.test:
cfg.data.train['ann_file'] = cfg.data.val['ann_file']
cfg.data.train['img_prefix'] = cfg.data.val['img_prefix']
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
num_gpus = args.gpus
rank = 0
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
num_gpus = torch.cuda.device_count()
rank, _ = get_dist_info()
if cfg.optimizer['type'] == 'SGD':
cfg.optimizer['lr'] *= num_gpus * cfg.data.imgs_per_gpu / 256
else:
cfg.optimizer['lr'] *= ((num_gpus / 8) * (cfg.data.imgs_per_gpu / 2))
# init logger before other steps
logger = get_root_logger(nlogger, cfg.log_level)
if rank == 0:
logger.set_logger_dir(cfg.work_dir, 'd')
logger.info("Config: ------------------------------------------\n" +
cfg.text)
logger.info('Distributed training: {}'.format(distributed))
# set random seeds
if args.seed is not None:
logger.info('Set random seed to {}'.format(args.seed))
set_random_seed(args.seed)
model = build_detector(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
if rank == 0:
# describe_vars(model)
writer = set_writer(cfg.work_dir)
# try:
# # describe_features(model.backbone)
# writer.add_graph(model, torch.zeros((1, 3, 800, 800)))
# except (NotImplementedError, TypeError):
# logger.warn("Add graph failed.")
# except Exception as e:
# logger.warn("Add graph failed:", e)
if not args.graph and not args.profiler and not args.speed:
if distributed:
model = MMDistributedDataParallel(model.cuda())
else:
model = MMDataParallel(model, device_ids=range(cfg.gpus)).cuda()
if isinstance(cfg.data.train, list):
for t in cfg.data.train:
logger.info("loading training set: " + str(t.ann_file))
train_dataset = [build_dataset(t) for t in cfg.data.train]
CLASSES = train_dataset[0].CLASSES
else:
logger.info("loading training set: " +
str(cfg.data.train.ann_file))
train_dataset = build_dataset(cfg.data.train)
logger.info("{} images loaded!".format(len(train_dataset)))
CLASSES = train_dataset.CLASSES
if cfg.checkpoint_config is not None:
# save mmdet version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(mmdet_version=__version__,
config=cfg.text,
CLASSES=CLASSES)
# add an attribute for visualization convenience
if hasattr(model, 'module'):
model.module.CLASSES = CLASSES
else:
model.CLASSES = CLASSES
train_detector(model,
train_dataset,
cfg,
distributed=distributed,
validate=args.validate,
logger=logger,
runner_attr_dict={'task_name': args.dir})
else:
from mmcv.runner.checkpoint import load_checkpoint
from mmdet.datasets import build_dataloader
from mmdet.core.utils.model_utils import register_hooks
from mmdet.apis.train import parse_losses
model = MMDataParallel(model, device_ids=range(cfg.gpus)).cuda()
if args.profiler == 'test' or args.speed == 'test':
model.eval()
dataset = build_dataset(cfg.data.test)
else:
model.train()
dataset = build_dataset(cfg.data.train)
if cfg.load_from and (args.profiler or args.speed):
logger.info('load checkpoint from %s', cfg.load_from)
load_checkpoint(model,
cfg.load_from,
map_location='cpu',
strict=True)
data_loader = build_dataloader(dataset,
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
cfg.gpus,
dist=False,
shuffle=False)
if args.graph:
id_dict = {}
for name, parameter in model.named_parameters():
id_dict[id(parameter)] = name
for i, data_batch in enumerate(data_loader):
if args.graph:
outputs = model(**data_batch)
loss, log_vars = parse_losses(outputs)
get_dot = register_hooks(loss, id_dict)
loss.backward()
dot = get_dot()
dot.save('graph.dot')
break
elif args.profiler:
with torch.autograd.profiler.profile(use_cuda=True) as prof:
if args.profiler == 'train':
outputs = model(**data_batch)
loss, log_vars = parse_losses(outputs)
loss.backward()
else:
with torch.no_grad():
model(**data_batch, return_loss=False)
if i == 20:
prof.export_chrome_trace('./trace.json')
logger.info(prof)
break
elif args.speed:
if args.speed == 'train':
start = time.perf_counter()
outputs = model(**data_batch)
loss, log_vars = parse_losses(outputs)
loss.backward()
torch.cuda.synchronize()
end = time.perf_counter()
else:
start = time.perf_counter()
with torch.no_grad():
model(**data_batch, return_loss=False)
end = time.perf_counter()
logger.info("{:.3f} s/iter, {:.1f} iters/s".format(
end - start, 1. / (end - start)))
