def train(
work_dir,
model_cfg,
loss_cfg,
dataset_cfg,
optimizer_cfg,
total_epochs,
training_hooks,
batch_size=None,
gpu_batch_size=None,
workflow=[('train', 1)],
gpus=-1,
log_level=0,
workers=4,
resume_from=None,
load_from=None,
):
# calculate batch size
if gpus < 0:
gpus = torch.cuda.device_count()
if (batch_size is None) and (gpu_batch_size is not None):
batch_size = gpu_batch_size * gpus
assert batch_size is not None, 'Please appoint batch_size or gpu_batch_size.'
# prepare data loaders
if isinstance(dataset_cfg, dict):
dataset_cfg = [dataset_cfg]
data_loaders = [
torch.utils.data.DataLoader(dataset=call_obj(**d),
batch_size=batch_size,
shuffle=True,
num_workers=workers,
drop_last=True) for d in dataset_cfg
]
# put model on gpus
if isinstance(model_cfg, list):
model = [call_obj(**c) for c in model_cfg]
model = torch.nn.Sequential(*model)
else:
model = call_obj(**model_cfg)
model.apply(weights_init)
model = MMDataParallel(model, device_ids=range(gpus)).cuda()
loss = call_obj(**loss_cfg)
# build runner
optimizer = call_obj(params=model.parameters(), **optimizer_cfg)
runner = Runner(model, batch_processor, optimizer, work_dir, log_level)
runner.register_training_hooks(**training_hooks)
if resume_from:
runner.resume(resume_from)
elif load_from:
runner.load_checkpoint(load_from)
# run
workflow = [tuple(w) for w in workflow]
runner.run(data_loaders, workflow, total_epochs, loss=loss)
def train(
work_dir,
model_cfg,
loss_cfg,
dataset_cfg,
optimizer_cfg,
batch_size,
total_epochs,
training_hooks,
workflow=[('train', 1)],
gpus=1,
log_level=0,
workers=2,
resume_from=None,
load_from=None,
):
# prepare data loaders
if isinstance(dataset_cfg, dict):
dataset_cfg = [dataset_cfg]
data_loaders = [
torch.utils.data.DataLoader(dataset=call_obj(**d),
batch_size=batch_size,
shuffle=True,
num_workers=workers,
drop_last=True) for d in dataset_cfg
]
# put model on gpus
if isinstance(model_cfg, list):
model = [call_obj(**c) for c in model_cfg]
model = torch.nn.Sequential(*model)
else:
model = call_obj(**model_cfg)
model.apply(weights_init)
print("Model size: ",
sum(p.numel() for p in model.parameters() if p.requires_grad))
model = MMDataParallel(model, device_ids=range(gpus)).cuda()
loss = call_obj(**loss_cfg)
# build runner
optimizer = call_obj(params=model.parameters(), **optimizer_cfg)
runner = Runner(model, batch_processor, optimizer, work_dir, log_level)
runner.register_training_hooks(**training_hooks)
if resume_from:
runner.resume(resume_from)
elif load_from:
runner.load_checkpoint(load_from)
# run
workflow = [tuple(w) for w in workflow]
runner.run(data_loaders, workflow, total_epochs, loss=loss)
writer.export_scalars_to_json("./all_scalars.json")
writer.close()
def train(
work_dir,
model_cfg,
dataset_cfg,
batch_size,
optimizer_cfg,
total_epochs,
training_hooks,
workflow=[('train', 1)],
gpus=1,
log_level=0,
workers=4,
resume_from=None,
load_from=None,
):
# prepare data loaders
if isinstance(dataset_cfg, dict):
dataset_cfg = [dataset_cfg]
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
data_loaders = [
torch.utils.data.DataLoader(
dataset=call_obj(**d,
transform=transforms.Compose([
transforms.ToTensor(),
normalize,
])),
batch_size=batch_size * gpus,
shuffle=True,
num_workers=workers,
drop_last=True) for d in dataset_cfg
]
# put model on gpus
if isinstance(model_cfg, list):
model = [call_obj(**c) for c in model_cfg]
model = torch.nn.Sequential(*model)
else:
model = call_obj(**model_cfg)
model = MMDataParallel(model, device_ids=range(gpus)).cuda()
# build runner
optimizer = call_obj(params=model.parameters(), **optimizer_cfg)
runner = Runner(model, batch_processor, optimizer, work_dir, log_level)
runner.register_training_hooks(**training_hooks)
if resume_from:
runner.resume(resume_from)
elif load_from:
runner.load_checkpoint(load_from)
# run
workflow = [tuple(w) for w in workflow]
runner.run(data_loaders, workflow, total_epochs)
def __init__(self, cfg):
super(TrackerCnnModule, self).__init__()
model = build_detector(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
checkpoint = glob.glob(os.path.join(mmdet_dir, work_dir, '*.pth'))[0]
load_checkpoint(model, checkpoint)
model = MMDataParallel(model, device_ids=[0])
self.bbox_roi_extractor = builder.build_roi_extractor(
cfg.bbox_roi_extractor)
for param in model.parameters():
param.requires_grad = False
self.detector = model.module
self.flatten = Flatten()
self.cfg = cfg
def train_flownet(model,
dataset,
cfg,
distributed=False,
validate=False,
logger=None):
if logger is None:
logger = get_root_logger(cfg.log_level)
# start training
# prepare data loaders
data_loaders = [
build_dataloader(dataset,
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
cfg.gpus,
dist=False)
]
# put model on gpus
model = MMDataParallel(model, device_ids=range(cfg.gpus)).cuda()
# build runner
runner = Runner(model, batch_processor, cfg.optimizer, cfg.work_dir,
cfg.log_level)
runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config,
cfg.checkpoint_config, cfg.log_config)
# if cfg.resume_from:
# runner.resume(cfg.resume_from)
# elif cfg.load_from:
# runner.load_checkpoint(cfg.load_from)
model.eval()
for param in model.parameters():
param.requires_grad = False
# model.load_flow()
model.module.flow_head.train()
for param in model.module.flow_head.parameters():
param.requires_grad = True
# training
runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
def main():
args = parse_args()
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.data.test.test_mode = True
if cfg.data.test.oversample == 'three_crop':
cfg.model.spatial_temporal_module.spatial_size = 8
dataset = obj_from_dict(cfg.data.test, datasets, dict(test_mode=True))
if args.gpus == 1:
model = build_recognizer(cfg.model,
train_cfg=None,
test_cfg=cfg.test_cfg)
load_checkpoint(model, args.checkpoint, strict=True)
model = MMDataParallel(model, device_ids=[1])
params = list(model.parameters())
weight_softmax = np.squeeze(params[-2].data.cpu().numpy(
)) # fully conneted layer parameters to numpy already
data_loader = build_dataloader(
dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
num_gpus=1,
dist=False,
shuffle=False)
outputs, inputs = single_test(model, data_loader)
else:
model_args = cfg.model.copy()
model_args.update(train_cfg=None, test_cfg=cfg.test_cfg)
model_type = getattr(recognizers, model_args.pop('type'))
outputs = parallel_test(model_type,
model_args,
args.checkpoint,
dataset,
_data_func,
range(args.gpus),
workers_per_gpu=args.proc_per_gpu)
#print(len(features_blobs))
#print(features_blobs[0].size())
if args.out:
print('writing results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
num_videos = len(outputs)
class_name = 'YoYo'
os.mkdir('data/CAM_imgs/' + class_name)
for k in range(0, num_videos):
os.mkdir('data/CAM_imgs/' + class_name + '/CAMs_{:02d}'.format(k))
idx = get_top_5_index("tools/results.pkl",
k) # change the dir of results.pkl to tools/
conv_feat = pickle.load(open(
"tools/hook_features/feat_{:02d}.pkl".format(k), 'rb'),
encoding='utf-8')
conv_feat = conv_feat.cpu().numpy()
CAMs = returnCAM(
conv_feat, weight_softmax,
[idx[0]
]) # generate class activation mapping for the top1 prediction
single_input = inputs[k].numpy()
writeCAMs(class_name, CAMs, single_input, k)
gt_labels = []
for i in range(len(dataset)):
ann = dataset.get_ann_info(i)
gt_labels.append(ann['label'])
if args.use_softmax:
print("Averaging score over {} clips with softmax".format(
outputs[0].shape[0]))
results = [softmax(res, dim=1).mean(axis=0) for res in outputs]
else:
print("Averaging score over {} clips without softmax (ie, raw)".format(
outputs[0].shape[0]))
results = [res.mean(axis=0) for res in outputs]
top1, top5 = top_k_accuracy(results, gt_labels, k=(1, 5))
mean_acc = mean_class_accuracy(results, gt_labels)
print("Mean Class Accuracy = {:.02f}".format(mean_acc * 100))
print("Top-1 Accuracy = {:.02f}".format(top1 * 100))
print("Top-5 Accuracy = {:.02f}".format(top5 * 100))
def _non_dist_train(model,
dataset,
cfg,
validate=False,
logger=None,
timestamp=None,
meta=None):
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = [
build_dataloader(ds,
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
len(cfg.gpu_ids),
dist=False,
seed=cfg.seed) for ds in dataset
]
# put model on gpus
model = MMDataParallel(model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)
# build runner
if cfg.optimizer.type == 'SGD_GC':
optimizer = SGD_GC(model.parameters(),
cfg.optimizer.lr,
momentum=cfg.optimizer.momentum,
weight_decay=cfg.optimizer.weight_decay)
else:
optimizer = build_optimizer(model, cfg.optimizer)
runner = Runner(model,
batch_processor,
optimizer,
cfg.work_dir,
logger=logger,
meta=meta)
# an ugly walkaround to make the .log and .log.json filenames the same
runner.timestamp = timestamp
# fp16 setting
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
optimizer_config = Fp16OptimizerHook(**cfg.optimizer_config,
**fp16_cfg,
distributed=False)
else:
optimizer_config = cfg.optimizer_config
runner.register_training_hooks(cfg.lr_config, optimizer_config,
cfg.checkpoint_config, cfg.log_config)
# register eval hooks
if validate:
val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
val_dataloader = build_dataloader(
val_dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=False,
shuffle=False)
eval_cfg = cfg.get('evaluation', {})
runner.register_hook(EvalHook(val_dataloader, **eval_cfg))
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
def _non_dist_train(model, dataset, cfg, validate=False):
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = []
# for idx, ds in enumerate(dataset):
# if idx==0:
# data_loaders.append(build_dataloader(
# ds,
# cfg.data.imgs_per_gpu,
# cfg.data.workers_per_gpu,
# cfg.gpus,
# dist=False,
# cls_balanced_sampler=False))
# else:
# data_loaders.append(build_dataloader(
# ds,
# cfg.data.imgs_per_gpu,
# cfg.data.workers_per_gpu,
# cfg.gpus,
# dist=False,
# cls_balanced_sampler=True))
data_loaders.append(
build_dataloader(dataset[0],
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
cfg.gpus,
dist=False,
cls_balanced_sampler=True))
# put model on gpus
model = MMDataParallel(model, device_ids=range(cfg.gpus)).cuda()
load_checkpoint(model, cfg.load_from)
print('load from {}'.format(cfg.load_from))
model.eval()
for param in model.parameters():
param.requires_grad = False
# for param in model.parameters():
# param.requires_grad = False
# model.eval()
dataset_to_longtailmodel = {}
dataset_to_longtailmodel['train'] = data_loaders[0]
# dataset_to_longtailmodel['train_plain'] = data_loaders[1]
# dataset_to_longtailmodel['val'] = data_loaders[2]
## which to use
import datetime
# cls_models = {
# 1: '2fc_rand',
# 2: '3fc_rand',
# 3: '3fc_ft' }
use_model = cfg.use_model
print('use {}'.format(cfg.use_model))
exp_prefix = cfg.exp_prefix
total_epoch = 120
initial_lr = 0.01
if hasattr(dataset[0], 'coco'):
if use_model == '2fc_rand':
cls_head = simple2fc(num_classes=81).cuda()
elif use_model == '3fc_rand' or '3fc_ft':
cls_head = simple3fc(num_classes=81).cuda()
elif hasattr(dataset[0], 'lvis'):
if use_model == '2fc_rand':
cls_head = simple2fc(num_classes=1231).cuda()
elif use_model == '3fc_rand' or '3fc_ft':
cls_head = simple3fc(num_classes=1231).cuda()
optimizer = optim.SGD([{
'params': cls_head.parameters(),
'lr': initial_lr
}])
# for param in list(cls_head.parameters())[:-4]:
# param.requires_grad = False
def save_ckpt(cls_head):
save_checkpoint(
cls_head, './{}/{}_{}.pth'.format(cfg.work_dir, exp_prefix,
use_model))
torch.save(
epoch, './{}/{}_{}_epoch.pth'.format(cfg.work_dir, exp_prefix,
use_model))
def load_ckpt(cls_head, use_model):
if use_model == '2fc_rand' or use_model == '3fc_rand':
if not os.path.exists('./{}_{}.pth'.format(exp_prefix, use_model)):
print('start training from 0 epoch')
return 0
else:
epoch = torch.load('./{}_{}_epoch.pth'.format(
exp_prefix, use_model))
load_checkpoint(cls_head,
'./{}_{}.pth'.format(exp_prefix, use_model))
return epoch
elif use_model == '3fc_ft':
if not os.path.exists('./{}_{}.pth'.format(exp_prefix, use_model)):
print(
'start training from 0 epoch, init from orig 3fc cls head')
orig_head_state_dict = model.module.bbox_head.state_dict()
key_map = {
'fc_cls.weight': 'feat_classifier.fc_classifier.weight',
'fc_cls.bias': 'feat_classifier.fc_classifier.bias',
'shared_fcs.0.weight': 'feat_classifier.fc1.weight',
'shared_fcs.0.bias': 'feat_classifier.fc1.bias',
'shared_fcs.1.weight': 'feat_classifier.fc2.weight',
'shared_fcs.1.bias': 'feat_classifier.fc2.bias'
}
new_state_dict = OrderedDict()
for key, value in orig_head_state_dict.items():
if key in key_map:
new_key = key_map[key]
new_state_dict[new_key] = value
cls_head.load_state_dict(new_state_dict)
return 0
else:
epoch = torch.load('./{}_{}_epoch.pth'.format(
exp_prefix, use_model))
load_checkpoint(cls_head,
'./{}_{}.pth'.format(exp_prefix, use_model))
return epoch
epoch = load_ckpt(cls_head, use_model)
cls_head = MMDataParallel(cls_head, device_ids=range(cfg.gpus)).cuda()
for epoch in range(epoch + 1, total_epoch + 1):
##due to schdualer bug, we do manual lr schedule
if epoch >= 8:
for param_group in optimizer.param_groups:
param_group['lr'] = initial_lr * 0.1
if epoch >= 11:
for param_group in optimizer.param_groups:
param_group['lr'] = initial_lr * 0.01
print('epoch {} lr {}'.format(epoch, optimizer.param_groups[0]['lr']))
for step, data_batch in enumerate(dataset_to_longtailmodel['train']):
if step % 10 == 0:
print('step {} time: {}'.format(step, datetime.datetime.now()))
torch.cuda.empty_cache()
bs = dataset_to_longtailmodel['train'].batch_size
nc = dataset_to_longtailmodel['train'].sampler.nc
data_cls_index = np.split(
dataset_to_longtailmodel['train'].sampler.class_indices,
dataset_to_longtailmodel['train'].sampler.episode, 0)[step]
neg_feats, gt_inds, gt_num, pos_feats, pos_label = model(
**data_batch)
samples_convert, label_converts = compose_training_data(
neg_feats, gt_inds, gt_num, pos_feats, pos_label,
data_cls_index, bs, nc)
logits = cls_head(samples_convert)
log_p_y = F.log_softmax(logits, dim=1).view(logits.size(0), -1)
loss_target = -log_p_y.gather(
1, label_converts.unsqueeze(1)).squeeze().view(-1).mean()
loss_val = 1.0 * loss_target
_, y_hat = log_p_y.max(1)
acc_val = torch.eq(y_hat, label_converts).float().mean()
bg_acc = (y_hat[label_converts == 0]
== 0).sum() / (label_converts == 0).sum().float()
fg_acc = torch.eq(
y_hat, label_converts).float()[label_converts != 0].mean()
if step % 10 == 0:
print('step {} acc: {}'.format(step, acc_val.item()))
print('step {} bg acc: {}'.format(step, bg_acc.item()))
print('step {} fg acc: {}'.format(step, fg_acc.item()))
optimizer.zero_grad()
loss_val.backward()
optimizer.step()
save_ckpt(cls_head)
def main():
args = parse_args()
cfg = Config.fromfile(args.cfg)
work_dir = cfg.work_dir
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
str(device_id) for device_id in cfg.device_ids)
log_dir = os.path.join(work_dir, 'logs')
if not os.path.exists(log_dir):
os.makedirs(log_dir)
logger = init_logger(log_dir)
seed = cfg.seed
logger.info('Set random seed to {}'.format(seed))
set_random_seed(seed)
train_dataset = get_dataset(cfg.data.train)
train_data_loader = build_dataloader(
train_dataset,
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
len(cfg.device_ids),
dist=False,
)
val_dataset = get_dataset(cfg.data.val)
val_data_loader = build_dataloader(val_dataset,
1,
cfg.data.workers_per_gpu,
1,
dist=False,
shuffle=False)
model = build_detector(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
model = MMDataParallel(model).cuda()
optimizer = obj_from_dict(cfg.optimizer, torch.optim,
dict(params=model.parameters()))
lr_scheduler = obj_from_dict(cfg.lr_scedule, LRschedule,
dict(optimizer=optimizer))
checkpoint_dir = os.path.join(cfg.work_dir, 'checkpoint_dir')
if not os.path.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
start_epoch = cfg.start_epoch
if cfg.resume_from:
checkpoint = load_checkpoint(model, cfg.resume_from)
start_epoch = 0
logger.info('resumed epoch {}, from {}'.format(start_epoch,
cfg.resume_from))
log_buffer = LogBuffer()
for epoch in range(start_epoch, cfg.end_epoch):
train(train_data_loader, model, optimizer, epoch, lr_scheduler,
log_buffer, cfg, logger)
tmp_checkpoint_file = os.path.join(checkpoint_dir, 'tmp_val.pth')
meta_dict = cfg._cfg_dict
logger.info('save tmp checkpoint to {}'.format(tmp_checkpoint_file))
save_checkpoint(model, tmp_checkpoint_file, optimizer, meta=meta_dict)
if len(cfg.device_ids) == 1:
sensitivity = val(val_data_loader, model, cfg, logger, epoch)
else:
model_args = cfg.model.copy()
model_args.update(train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
model_type = getattr(detectors, model_args.pop('type'))
results = parallel_test(
cfg,
model_type,
model_args,
tmp_checkpoint_file,
val_dataset,
np.arange(len(cfg.device_ids)).tolist(),
workers_per_gpu=1,
)
sensitivity = evaluate_deep_lesion(results, val_dataset,
cfg.cfg_3dce, logger)
save_file = os.path.join(
checkpoint_dir, 'epoch_{}_sens@4FP_{:.5f}_{}.pth'.format(
epoch + 1, sensitivity,
time.strftime('%m-%d-%H-%M', time.localtime(time.time()))))
os.rename(tmp_checkpoint_file, save_file)
logger.info('save checkpoint to {}'.format(save_file))
if epoch > cfg.lr_scedule.T_max:
os.remove(save_file)
