def main(args):
utils.init_distributed_mode(args)
print("git:\n {}\n".format(utils.get_sha()))
if args.frozen_weights is not None:
assert args.masks, "Frozen training is meant for segmentation only"
print(args)
# Save our Wandb metadata
if not args.no_wb:
wandb.init(entity='dl-project',
project='dl-final-project',
name=args.wb_name,
notes=args.wb_notes,
reinit=True)
wandb.config.epochs = args.epochs
device = torch.device(args.device)
# fix the seed for reproducibility
seed = args.seed + utils.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
model, criterion, postprocessors = build_model(args)
model.to(device)
# visualize_video(model, postprocessors)
model_without_ddp = model
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of trainable params:', n_parameters)
wandb.config.n_parameters = n_parameters
wandb.config.n_trainable_parameters = n_parameters # better name
# Log total # of model parameters (including frozen) to W&B
n_total_parameters = sum(p.numel() for p in model.parameters())
print('total number of parameters:', n_total_parameters)
wandb.config.n_total_parameters = n_total_parameters
dataset_train = build_dataset(image_set='train', args=args)
dataset_val = build_dataset(image_set='val', args=args)
# For visualization we want the raw images without any normalization or random resizing
dataset_val_without_resize = CocoDetection(
"data/coco/val2017",
annFile="data/coco/annotations/instances_val2017.json",
transforms=T.Compose([T.ToTensor()]))
# Save metadata about training + val datasets and batch size
wandb.config.len_dataset_train = len(dataset_train)
wandb.config.len_dataset_val = len(dataset_val)
wandb.config.batch_size = args.batch_size
if args.distributed:
if args.cache_mode:
sampler_train = samplers.NodeDistributedSampler(dataset_train)
sampler_val = samplers.NodeDistributedSampler(dataset_val,
shuffle=False)
else:
sampler_train = samplers.DistributedSampler(dataset_train)
sampler_val = samplers.DistributedSampler(dataset_val,
shuffle=False)
else:
sampler_train = torch.utils.data.RandomSampler(dataset_train)
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
batch_sampler_train = torch.utils.data.BatchSampler(sampler_train,
args.batch_size,
drop_last=True)
data_loader_train = DataLoader(dataset_train,
batch_sampler=batch_sampler_train,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
data_loader_val = DataLoader(dataset_val,
args.batch_size,
sampler=sampler_val,
drop_last=False,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
# lr_backbone_names = ["backbone.0", "backbone.neck", "input_proj", "transformer.encoder"]
def match_name_keywords(n, name_keywords):
out = False
for b in name_keywords:
if b in n:
out = True
break
return out
for n, p in model_without_ddp.named_parameters():
print(n)
param_dicts = [{
"params": [
p for n, p in model_without_ddp.named_parameters()
if not match_name_keywords(n, args.lr_backbone_names)
and not match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters() if
match_name_keywords(n, args.lr_backbone_names) and p.requires_grad
],
"lr":
args.lr_backbone,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters()
if match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr * args.lr_linear_proj_mult,
}]
# Not sure if we should save all hyperparameters in wandb.config?
# just start with a few important ones
wandb.config.lr = args.lr
wandb.config.lr_backbone = args.lr_backbone
if args.sgd:
optimizer = torch.optim.SGD(param_dicts,
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.AdamW(param_dicts,
lr=args.lr,
weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
if args.dataset_file == "coco_panoptic":
# We also evaluate AP during panoptic training, on original coco DS
coco_val = datasets.coco.build("val", args)
base_ds = get_coco_api_from_dataset(coco_val)
else:
base_ds = get_coco_api_from_dataset(dataset_val)
if args.frozen_weights is not None:
checkpoint = torch.load(args.frozen_weights, map_location='cpu')
model_without_ddp.detr.load_state_dict(checkpoint['model'])
output_dir = Path(args.output_dir)
if args.resume:
if args.resume.startswith('https'):
checkpoint = torch.hub.load_state_dict_from_url(args.resume,
map_location='cpu',
check_hash=True)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
missing_keys, unexpected_keys = model_without_ddp.load_state_dict(
checkpoint['model'], strict=False)
unexpected_keys = [
k for k in unexpected_keys
if not (k.endswith('total_params') or k.endswith('total_ops'))
]
if len(missing_keys) > 0:
print('Missing Keys: {}'.format(missing_keys))
if len(unexpected_keys) > 0:
print('Unexpected Keys: {}'.format(unexpected_keys))
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
import copy
p_groups = copy.deepcopy(optimizer.param_groups)
optimizer.load_state_dict(checkpoint['optimizer'])
for pg, pg_old in zip(optimizer.param_groups, p_groups):
pg['lr'] = pg_old['lr']
pg['initial_lr'] = pg_old['initial_lr']
# print(optimizer.param_groups)
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
# todo: this is a hack for doing experiment that resume from checkpoint and also modify lr scheduler (e.g., decrease lr in advance).
args.override_resumed_lr_drop = True
if args.override_resumed_lr_drop:
print(
'Warning: (hack) args.override_resumed_lr_drop is set to True, so args.lr_drop would override lr_drop in resumed lr_scheduler.'
)
lr_scheduler.step_size = args.lr_drop
lr_scheduler.base_lrs = list(
map(lambda group: group['initial_lr'],
optimizer.param_groups))
lr_scheduler.step(lr_scheduler.last_epoch)
args.start_epoch = checkpoint['epoch'] + 1
# check the resumed model
if not args.eval:
test_stats, coco_evaluator = evaluate(model, criterion,
postprocessors,
data_loader_val, base_ds,
device, args.output_dir)
if args.eval:
print("Generating visualizations...")
visualize_bbox(model, postprocessors, data_loader_val, device,
dataset_val_without_resize)
test_stats, coco_evaluator = evaluate(model, criterion, postprocessors,
data_loader_val, base_ds, device,
args.output_dir)
if args.output_dir:
utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval,
output_dir / "eval.pth")
return
print("Start training")
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
sampler_train.set_epoch(epoch)
train_stats = train_one_epoch(model, criterion, data_loader_train,
optimizer, device, epoch,
args.clip_max_norm)
lr_scheduler.step()
if args.output_dir:
checkpoint_file_for_wb = str(
output_dir / f'{wandb.run.id}_checkpoint{epoch:04}.pth')
checkpoint_paths = [
output_dir / 'checkpoint.pth', checkpoint_file_for_wb
]
# extra checkpoint before LR drop and every 5 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 5 == 0:
checkpoint_paths.append(output_dir /
f'checkpoint{epoch:04}.pth')
for checkpoint_path in checkpoint_paths:
utils.save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'args': args,
}, checkpoint_path)
# Save model checkpoint to W&B
wandb.save(checkpoint_file_for_wb)
# Generate visualizations for fixed(?) set of images every epoch
print("Generating visualizations...")
visualize_bbox(model, postprocessors, data_loader_val, device,
dataset_val_without_resize)
test_stats, coco_evaluator = evaluate(model, criterion, postprocessors,
data_loader_val, base_ds, device,
args.output_dir)
log_stats = {
**{f'train_{k}': v
for k, v in train_stats.items()},
**{f'test_{k}': v
for k, v in test_stats.items()}, 'epoch': epoch,
'n_parameters': n_parameters
}
# Save the COCO metrics properly
metric_name = [
"AP", "AP50", "AP75", "APs", "APm", "APl", "AR@1", "AR@10",
"AR@100", "ARs", "ARm", "ARl"
]
for i, metric_val in enumerate(log_stats["test_coco_eval_bbox"]):
log_stats[metric_name[i]] = metric_val
if not args.no_wb:
wandb.log(log_stats)
print("train_loss: ", log_stats['train_loss'])
if args.output_dir and utils.is_main_process():
with (output_dir / "log.txt").open("a") as f:
f.write(json.dumps(log_stats) + "\n")
wandb.save(str(output_dir / "log.txt"))
# for evaluation logs
if coco_evaluator is not None:
(output_dir / 'eval').mkdir(exist_ok=True)
if "bbox" in coco_evaluator.coco_eval:
filenames = ['latest.pth']
eval_filename_for_wb = f'{wandb.run.id}_eval_{epoch:04}.pth'
eval_path_for_wb = str(output_dir / "eval" /
eval_filename_for_wb)
filenames = ['latest.pth', eval_filename_for_wb]
if epoch % 50 == 0:
filenames.append(f'{epoch:03}.pth')
for name in filenames:
torch.save(coco_evaluator.coco_eval["bbox"].eval,
output_dir / "eval" / name)
# TODO not sure if this file will end up being too big
# I think it's the COCO precision/recall metrics
# in some format...
# let's track it just in case to start!
wandb.save(eval_path_for_wb)
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
def main(args):
utils.init_distributed_mode(args)
print("git:\n {}\n".format(utils.get_sha()))
if args.frozen_weights is not None:
assert args.masks, "Frozen training is meant for segmentation only"
print(args)
device = torch.device(args.device)
# fix the seed for reproducibility
seed = args.seed + utils.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
model, criterion, postprocessors = build_model(args)
model.to(device)
model_without_ddp = model
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
dataset_train = build_dataset(image_set='train', args=args)
dataset_val = build_dataset(image_set='val', args=args)
dataset_test = build_dataset(image_set='test', args=args)
if args.distributed:
if args.cache_mode:
sampler_train = samplers.NodeDistributedSampler(dataset_train)
sampler_val = samplers.NodeDistributedSampler(dataset_val,
shuffle=False)
else:
sampler_train = samplers.DistributedSampler(dataset_train)
sampler_val = samplers.DistributedSampler(dataset_val,
shuffle=False)
sampler_test = samplers.DistributedSampler(dataset_test,
shuffle=False)
else:
sampler_train = torch.utils.data.RandomSampler(dataset_train)
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
sampler_test = torch.utils.data.SequentialSampler(dataset_test)
batch_sampler_train = torch.utils.data.BatchSampler(sampler_train,
args.batch_size,
drop_last=True)
data_loader_train = DataLoader(dataset_train,
batch_sampler=batch_sampler_train,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
data_loader_val = DataLoader(dataset_val,
args.batch_size,
sampler=sampler_val,
drop_last=False,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
data_loader_test = DataLoader(dataset_test,
args.batch_size,
sampler=sampler_val,
drop_last=False,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
# lr_backbone_names = ["backbone.0", "backbone.neck", "input_proj", "transformer.encoder"]
def match_name_keywords(n, name_keywords):
out = False
for b in name_keywords:
if b in n:
out = True
break
return out
for n, p in model_without_ddp.named_parameters():
print(n)
param_dicts = [{
"params": [
p for n, p in model_without_ddp.named_parameters()
if not match_name_keywords(n, args.lr_backbone_names)
and not match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters() if
match_name_keywords(n, args.lr_backbone_names) and p.requires_grad
],
"lr":
args.lr_backbone,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters()
if match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr * args.lr_linear_proj_mult,
}]
if args.sgd:
optimizer = torch.optim.SGD(param_dicts,
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.AdamW(param_dicts,
lr=args.lr,
weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
if args.dataset_file == "coco_panoptic":
# We also evaluate AP during panoptic training, on original coco DS
coco_val = datasets.coco.build("val", args)
base_ds = get_coco_api_from_dataset(coco_val)
else:
base_ds = get_coco_api_from_dataset(dataset_val)
if args.frozen_weights is not None:
checkpoint = torch.load(args.frozen_weights, map_location='cpu')
model_without_ddp.detr.load_state_dict(checkpoint['model'])
output_dir = Path(args.output_dir)
if args.resume:
if args.resume.startswith('https'):
checkpoint = torch.hub.load_state_dict_from_url(args.resume,
map_location='cpu',
check_hash=True)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
del checkpoint["model"]["transformer.decoder.class_embed.0.weight"]
del checkpoint["model"]["transformer.decoder.class_embed.0.bias"]
del checkpoint["model"]["transformer.decoder.class_embed.1.weight"]
del checkpoint["model"]["transformer.decoder.class_embed.1.bias"]
del checkpoint["model"]["transformer.decoder.class_embed.2.weight"]
del checkpoint["model"]["transformer.decoder.class_embed.2.bias"]
del checkpoint["model"]["transformer.decoder.class_embed.3.weight"]
del checkpoint["model"]["transformer.decoder.class_embed.3.bias"]
del checkpoint["model"]["transformer.decoder.class_embed.4.weight"]
del checkpoint["model"]["transformer.decoder.class_embed.4.bias"]
del checkpoint["model"]["transformer.decoder.class_embed.5.weight"]
del checkpoint["model"]["transformer.decoder.class_embed.5.bias"]
del checkpoint["model"]["transformer.decoder.class_embed.6.weight"]
del checkpoint["model"]["transformer.decoder.class_embed.6.bias"]
del checkpoint["model"]["class_embed.0.weight"]
del checkpoint["model"]["class_embed.0.bias"]
del checkpoint["model"]["class_embed.1.weight"]
del checkpoint["model"]["class_embed.1.bias"]
del checkpoint["model"]["class_embed.2.weight"]
del checkpoint["model"]["class_embed.2.bias"]
del checkpoint["model"]["class_embed.3.weight"]
del checkpoint["model"]["class_embed.3.bias"]
del checkpoint["model"]["class_embed.4.weight"]
del checkpoint["model"]["class_embed.4.bias"]
del checkpoint["model"]["class_embed.5.weight"]
del checkpoint["model"]["class_embed.5.bias"]
del checkpoint["model"]["class_embed.6.weight"]
del checkpoint["model"]["class_embed.6.bias"]
missing_keys, unexpected_keys = model_without_ddp.load_state_dict(
checkpoint['model'], strict=False)
unexpected_keys = [
k for k in unexpected_keys
if not (k.endswith('total_params') or k.endswith('total_ops'))
]
# if len(missing_keys) > 0:
# print('Missing Keys: {}'.format(missing_keys))
# if len(unexpected_keys) > 0:
# print('Unexpected Keys: {}'.format(unexpected_keys))
# if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
# import copy
# p_groups = copy.deepcopy(optimizer.param_groups)
# optimizer.load_state_dict(checkpoint['optimizer'])
# for pg, pg_old in zip(optimizer.param_groups, p_groups):
# pg['lr'] = pg_old['lr']
# pg['initial_lr'] = pg_old['initial_lr']
# #print(optimizer.param_groups)
# lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
# # todo: this is a hack for doing experiment that resume from checkpoint and also modify lr scheduler (e.g., decrease lr in advance).
# args.override_resumed_lr_drop = True
# if args.override_resumed_lr_drop:
# print('Warning: (hack) args.override_resumed_lr_drop is set to True, so args.lr_drop would override lr_drop in resumed lr_scheduler.')
# lr_scheduler.step_size = args.lr_drop
# lr_scheduler.base_lrs = list(map(lambda group: group['initial_lr'], optimizer.param_groups))
# lr_scheduler.step(lr_scheduler.last_epoch)
# args.start_epoch = checkpoint['epoch'] + 1
# # check the resumed model
if not args.eval:
test_stats, coco_evaluator = evaluate(model, criterion,
postprocessors,
data_loader_val, base_ds,
device, args.output_dir)
if args.eval:
test_stats, coco_evaluator = evaluate(model, criterion, postprocessors,
data_loader_val, base_ds, device,
args.output_dir)
if args.output_dir:
utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval,
output_dir / "eval.pth")
return
if args.test:
test_stats, coco_evaluator = evaluate(model, criterion, postprocessors,
data_loader_test, base_ds,
device, args.output_dir)
if args.output_dir:
utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval,
output_dir / "eval.pth")
return
print("Start training")
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
sampler_train.set_epoch(epoch)
train_stats = train_one_epoch(model, criterion, data_loader_train,
optimizer, device, epoch,
args.clip_max_norm)
lr_scheduler.step()
if args.output_dir:
checkpoint_paths = [output_dir / 'checkpoint.pth']
# extra checkpoint before LR drop and every 5 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 5 == 0:
checkpoint_paths.append(output_dir /
f'checkpoint{epoch:04}.pth')
for checkpoint_path in checkpoint_paths:
utils.save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'args': args,
}, checkpoint_path)
test_stats, coco_evaluator = evaluate(model, criterion, postprocessors,
data_loader_val, base_ds, device,
args.output_dir)
log_stats = {
**{f'train_{k}': v
for k, v in train_stats.items()},
**{f'test_{k}': v
for k, v in test_stats.items()}, 'epoch': epoch,
'n_parameters': n_parameters
}
if args.output_dir and utils.is_main_process():
with (output_dir / "log.txt").open("a") as f:
f.write(json.dumps(log_stats) + "\n")
# for evaluation logs
if coco_evaluator is not None:
(output_dir / 'eval').mkdir(exist_ok=True)
if "bbox" in coco_evaluator.coco_eval:
filenames = ['latest.pth']
if epoch % 50 == 0:
filenames.append(f'{epoch:03}.pth')
for name in filenames:
torch.save(coco_evaluator.coco_eval["bbox"].eval,
output_dir / "eval" / name)
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
def main(args):
utils.init_distributed_mode(args)
print("git:\n {}\n".format(utils.get_sha()))
if args.frozen_weights is not None:
assert args.masks, "Frozen training is meant for segmentation only"
print(args)
device = torch.device(args.device)
# fix the seed for reproducibility
seed = args.seed + utils.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
scaler = torch.cuda.amp.GradScaler(enabled=args.fp16)
if args.det_val:
assert args.eval, 'only support eval mode of detector for track'
model, criterion, postprocessors = build_model(args)
elif args.eval:
model, criterion, postprocessors = build_tracktest_model(args)
else:
model, criterion, postprocessors = build_tracktrain_model(args)
model.to(device)
model_without_ddp = model
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
dataset_train = build_dataset(image_set=args.track_train_split, args=args)
dataset_val = build_dataset(image_set=args.track_eval_split, args=args)
if args.distributed:
if args.cache_mode:
sampler_train = samplers.NodeDistributedSampler(dataset_train)
sampler_val = samplers.NodeDistributedSampler(dataset_val,
shuffle=False)
else:
sampler_train = samplers.DistributedSampler(dataset_train)
sampler_val = samplers.DistributedSampler(dataset_val,
shuffle=False)
else:
sampler_train = torch.utils.data.RandomSampler(dataset_train)
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
batch_sampler_train = torch.utils.data.BatchSampler(sampler_train,
args.batch_size,
drop_last=True)
data_loader_train = DataLoader(dataset_train,
batch_sampler=batch_sampler_train,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
data_loader_val = DataLoader(dataset_val,
args.batch_size,
sampler=sampler_val,
drop_last=False,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
# lr_backbone_names = ["backbone.0", "backbone.neck", "input_proj", "transformer.encoder"]
def match_name_keywords(n, name_keywords):
out = False
for b in name_keywords:
if b in n:
out = True
break
return out
for n, p in model_without_ddp.named_parameters():
print(n)
param_dicts = [{
"params": [
p for n, p in model_without_ddp.named_parameters()
if not match_name_keywords(n, args.lr_backbone_names)
and not match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters() if
match_name_keywords(n, args.lr_backbone_names) and p.requires_grad
],
"lr":
args.lr_backbone,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters()
if match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr * args.lr_linear_proj_mult,
}]
if args.sgd:
optimizer = torch.optim.SGD(param_dicts,
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.AdamW(param_dicts,
lr=args.lr,
weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu], find_unused_parameters=True)
model_without_ddp = model.module
if args.dataset_file == "coco_panoptic":
# We also evaluate AP during panoptic training, on original coco DS
coco_val = datasets.coco.build("val", args)
base_ds = get_coco_api_from_dataset(coco_val)
else:
base_ds = get_coco_api_from_dataset(dataset_val)
if args.frozen_weights is not None:
checkpoint = torch.load(args.frozen_weights, map_location='cpu')
model_without_ddp.detr.load_state_dict(checkpoint['model'])
output_dir = Path(args.output_dir)
if args.resume:
if args.resume.startswith('https'):
checkpoint = torch.hub.load_state_dict_from_url(args.resume,
map_location='cpu',
check_hash=True)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
missing_keys, unexpected_keys = model_without_ddp.load_state_dict(
checkpoint['model'], strict=False)
unexpected_keys = [
k for k in unexpected_keys
if not (k.endswith('total_params') or k.endswith('total_ops'))
]
if len(missing_keys) > 0:
print('Missing Keys: {}'.format(missing_keys))
if len(unexpected_keys) > 0:
print('Unexpected Keys: {}'.format(unexpected_keys))
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
import copy
p_groups = copy.deepcopy(optimizer.param_groups)
optimizer.load_state_dict(checkpoint['optimizer'])
for pg, pg_old in zip(optimizer.param_groups, p_groups):
pg['lr'] = pg_old['lr']
pg['initial_lr'] = pg_old['initial_lr']
print(optimizer.param_groups)
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
# todo: this is a hack for doing experiment that resume from checkpoint and also modify lr scheduler (e.g., decrease lr in advance).
args.override_resumed_lr_drop = True
if args.override_resumed_lr_drop:
print(
'Warning: (hack) args.override_resumed_lr_drop is set to True, so args.lr_drop would override lr_drop in resumed lr_scheduler.'
)
lr_scheduler.step_size = args.lr_drop
lr_scheduler.base_lrs = list(
map(lambda group: group['initial_lr'],
optimizer.param_groups))
lr_scheduler.step(lr_scheduler.last_epoch)
args.start_epoch = checkpoint['epoch'] + 1
# check the resumed model
# if not args.eval:
# test_stats, coco_evaluator, _ = evaluate(
# model, criterion, postprocessors, data_loader_val, base_ds, device, args.output_dir
# )
if args.eval:
assert args.batch_size == 1, print("Now only support 1.")
tracker = Tracker(score_thresh=args.track_thresh)
test_stats, coco_evaluator, res_tracks = evaluate(model,
criterion,
postprocessors,
data_loader_val,
base_ds,
device,
args.output_dir,
tracker=tracker,
phase='eval',
det_val=args.det_val,
fp16=args.fp16)
if args.output_dir:
# utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval, output_dir / "eval.pth")
if res_tracks is not None:
print("Creating video index for {}.".format(args.dataset_file))
video_to_images = defaultdict(list)
video_names = defaultdict()
for _, info in dataset_val.coco.imgs.items():
video_to_images[info["video_id"]].append({
"image_id":
info["id"],
"frame_id":
info["frame_id"]
})
video_name = info["file_name"].split("/")[0]
if video_name not in video_names:
video_names[info["video_id"]] = video_name
assert len(video_to_images) == len(video_names)
# save mot results.
save_track(res_tracks, args.output_dir, video_to_images,
video_names, args.track_eval_split)
return
print("Start training")
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
sampler_train.set_epoch(epoch)
train_stats = train_one_epoch(model,
criterion,
data_loader_train,
optimizer,
device,
scaler,
epoch,
args.clip_max_norm,
fp16=args.fp16)
lr_scheduler.step()
if args.output_dir:
checkpoint_paths = [output_dir / 'checkpoint.pth']
# extra checkpoint before LR drop and every 5 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 5 == 0:
checkpoint_paths.append(output_dir /
f'checkpoint{epoch:04}.pth')
for checkpoint_path in checkpoint_paths:
utils.save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'args': args,
}, checkpoint_path)
log_stats = {
**{f'train_{k}': v
for k, v in train_stats.items()}, 'epoch': epoch,
'n_parameters': n_parameters
}
if epoch % 10 == 0 or epoch > args.epochs - 5:
test_stats, coco_evaluator, _ = evaluate(model,
criterion,
postprocessors,
data_loader_val,
base_ds,
device,
args.output_dir,
fp16=args.fp16)
log_test_stats = {
**{f'test_{k}': v
for k, v in test_stats.items()}
}
log_stats.update(log_test_stats)
if args.output_dir and utils.is_main_process():
with (output_dir / "log.txt").open("a") as f:
f.write(json.dumps(log_stats) + "\n")
# for evaluation logs
# if coco_evaluator is not None:
# (output_dir / 'eval').mkdir(exist_ok=True)
# if "bbox" in coco_evaluator.coco_eval:
# filenames = ['latest.pth']
# if epoch % 50 == 0:
# filenames.append(f'{epoch:03}.pth')
# for name in filenames:
# torch.save(coco_evaluator.coco_eval["bbox"].eval,
# output_dir / "eval" / name)
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
def main(args):
utils.init_distributed_mode(args)
print("git:\n {}\n".format(utils.get_sha()))
if args.frozen_weights is not None:
assert args.masks, "Frozen training is meant for segmentation only"
print(args)
device = torch.device(args.device)
# fix the seed for reproducibility
seed = args.seed + utils.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
Dataset = get_dataset(args.dataset, args.task)
f = open(args.data_cfg)
data_config = json.load(f)
trainset_paths = data_config['train']
dataset_root = data_config['root']
f.close()
normalize = T.Compose([
T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
scales = [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]
transforms = T.Compose([
T.RandomHorizontalFlip(),
T.RandomSelect(
T.RandomResize(scales, max_size=1333),
T.Compose([
T.RandomResize([400, 500, 600]),
T.RandomSizeCrop(384, 600),
# T.RandomSizeCrop_MOT(384, 600),
T.RandomResize(scales, max_size=1333),
])),
normalize,
])
dataset_train = Dataset(args,
dataset_root,
trainset_paths, (1088, 608),
augment=True,
transforms=transforms)
args.nID = dataset_train.nID
model, criterion, postprocessors = build_model(args)
model.to(device)
model_without_ddp = model
# dataset_train = build_dataset(image_set='train', args=args)
# dataset_val = build_dataset(image_set='val', args=args)
if args.distributed:
if args.cache_mode:
sampler_train = samplers.NodeDistributedSampler(dataset_train)
# sampler_val = samplers.NodeDistributedSampler(dataset_val, shuffle=False)
else:
sampler_train = samplers.DistributedSampler(dataset_train)
# sampler_val = samplers.DistributedSampler(dataset_val, shuffle=False)
else:
sampler_train = torch.utils.data.RandomSampler(dataset_train)
# sampler_val = torch.utils.data.SequentialSampler(dataset_val)
batch_sampler_train = torch.utils.data.BatchSampler(sampler_train,
args.batch_size,
drop_last=True)
data_loader_train = DataLoader(dataset_train,
batch_sampler=batch_sampler_train,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
# data_loader_val = DataLoader(dataset_val, args.batch_size, sampler=sampler_val,
# drop_last=False, collate_fn=utils.collate_fn, num_workers=args.num_workers,
# pin_memory=True)
# data_loader_train = torch.utils.data.DataLoader(
# dataset_train,
# batch_size=args.batch_size,
# shuffle=True,
# num_workers=args.num_workers,
# pin_memory=True,
# drop_last=True
# )
# lr_backbone_names = ["backbone.0", "backbone.neck", "input_proj", "transformer.encoder"]
def match_name_keywords(n, name_keywords):
out = False
for b in name_keywords:
if b in n:
out = True
break
return out
for n, p in model_without_ddp.named_parameters():
print(n)
# 用于将classifer不更新参数
# for name,p in model_without_ddp.named_parameters():
# if name.startswith('classifier'):
# p.requires_grad = False
param_dicts = [{
"params": [
p for n, p in model_without_ddp.named_parameters()
if not match_name_keywords(n, args.lr_backbone_names)
and not match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters() if
match_name_keywords(n, args.lr_backbone_names) and p.requires_grad
],
"lr":
args.lr_backbone,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters()
if match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr * args.lr_linear_proj_mult,
}]
if args.sgd:
optimizer = torch.optim.SGD(param_dicts,
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.AdamW(param_dicts,
lr=args.lr,
weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)
# optimizer.add_param_group({'params': criterion.parameters()})
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
if args.frozen_weights is not None:
checkpoint = torch.load(args.frozen_weights, map_location='cpu')
model_without_ddp.detr.load_state_dict(checkpoint['model'])
output_dir = Path(args.output_dir)
if args.resume:
if args.resume.startswith('https'):
checkpoint = torch.hub.load_state_dict_from_url(args.resume,
map_location='cpu',
check_hash=True)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
model_dict = model_without_ddp.state_dict() #当前模型参数
pretrained_dict = {
k: v
for k, v in checkpoint['model'].items() if k not in [
"class_embed.0.weight", "class_embed.0.bias",
"class_embed.1.weight", "class_embed.1.bias",
"class_embed.2.weight", "class_embed.2.bias",
"class_embed.3.weight", "class_embed.3.bias",
"class_embed.4.weight", "class_embed.4.bias",
"class_embed.5.weight", "class_embed.5.bias"
]
}
model_dict.update(pretrained_dict)
# missing_keys, unexpected_keys = model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
missing_keys, unexpected_keys = model_without_ddp.load_state_dict(
model_dict, strict=False)
unexpected_keys = [
k for k in unexpected_keys
if not (k.endswith('total_params') or k.endswith('total_ops'))
]
if len(missing_keys) > 0:
print('Missing Keys: {}'.format(missing_keys))
if len(unexpected_keys) > 0:
print('Unexpected Keys: {}'.format(unexpected_keys))
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
args.start_epoch = checkpoint['epoch'] + 1
# optimizer.load_state_dict(checkpoint['optimizer'])
# if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
# import copy
# p_groups = copy.deepcopy(optimizer.param_groups)
# # optimizer.load_state_dict(checkpoint['optimizer'])
# for pg, pg_old in zip(optimizer.param_groups, p_groups):
# pg['lr'] = pg_old['lr']
# pg['initial_lr'] = pg_old['initial_lr']
# # print(optimizer.param_groups)
# lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
# # todo: this is a hack for doing experiment that resume from checkpoint and also modify lr scheduler (e.g., decrease lr in advance).
# args.override_resumed_lr_drop = True
# if args.override_resumed_lr_drop:
# print('Warning: (hack) args.override_resumed_lr_drop is set to True, so args.lr_drop would override lr_drop in resumed lr_scheduler.')
# lr_scheduler.step_size = args.lr_drop
# lr_scheduler.base_lrs = list(map(lambda group: group['initial_lr'], optimizer.param_groups))
# lr_scheduler.step(lr_scheduler.last_epoch)
# model.add_module('id')
# [p for p in model.named_parameters() if not p[1].requires_grad]
# 用于将classifer不更新参数
# optimizer = torch.optim.SGD(filter(lambda x: "classifier" not in x[0], model.parameters()), lr=args.lr,
# momentum=0.9, weight_decay=1e-4)
# model.classifier.training = False
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
print("Start training")
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
sampler_train.set_epoch(epoch)
train_stats = train_one_epoch(args, model, criterion,
data_loader_train, optimizer, device,
epoch, args.clip_max_norm)
lr_scheduler.step()
if args.output_dir:
checkpoint_paths = [output_dir / 'checkpoint.pth']
# extra checkpoint before LR drop and every 5 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 5 == 0:
checkpoint_paths.append(output_dir /
f'checkpoint{epoch:04}.pth')
for checkpoint_path in checkpoint_paths:
utils.save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'args': args,
}, checkpoint_path)
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
def main(args):
utils.init_distributed_mode(args)
print(args)
device = torch.device(args.device)
# fix the seed for reproducibility
seed = args.seed + utils.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
model, criterion, postprocessors = build_model(args)
model.to(device)
model_without_ddp = model
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
image_set = 'fewshot' if args.fewshot_finetune else 'train'
dataset_train = build_dataset(image_set=image_set, args=args)
dataset_val = build_dataset(image_set='val', args=args)
dataset_support = build_support_dataset(image_set=image_set, args=args)
if args.distributed:
if args.cache_mode:
sampler_train = samplers.NodeDistributedSampler(dataset_train)
sampler_val = samplers.NodeDistributedSampler(dataset_val,
shuffle=False)
sampler_support = samplers.NodeDistributedSampler(dataset_support)
else:
sampler_train = samplers.DistributedSampler(dataset_train)
sampler_val = samplers.DistributedSampler(dataset_val,
shuffle=False)
sampler_support = samplers.DistributedSampler(dataset_support)
else:
sampler_train = torch.utils.data.RandomSampler(dataset_train)
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
sampler_support = torch.utils.data.RandomSampler(dataset_support)
batch_sampler_train = torch.utils.data.BatchSampler(sampler_train,
args.batch_size,
drop_last=False)
loader_train = DataLoader(dataset_train,
batch_sampler=batch_sampler_train,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
loader_val = DataLoader(dataset_val,
batch_size=args.batch_size,
sampler=sampler_val,
drop_last=False,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
loader_support = DataLoader(dataset_support,
batch_size=1,
sampler=sampler_support,
drop_last=False,
num_workers=args.num_workers,
pin_memory=False)
def match_name_keywords(n, name_keywords):
out = False
for b in name_keywords:
if b in n:
out = True
break
return out
for n, p in model_without_ddp.named_parameters():
print(n)
if not args.fewshot_finetune:
param_dicts = [{
"params": [
p for n, p in model_without_ddp.named_parameters()
if not match_name_keywords(n, args.lr_backbone_names)
and not match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr,
"initial_lr":
args.lr,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters()
if match_name_keywords(n, args.lr_backbone_names)
and p.requires_grad
],
"lr":
args.lr_backbone,
"initial_lr":
args.lr_backbone,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters()
if match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr * args.lr_linear_proj_mult,
"initial_lr":
args.lr * args.lr_linear_proj_mult,
}]
else:
# For few-shot finetune stage, do not train sampling offsets, reference points, and embedding related parameters
param_dicts = [
{
"params":
[p for n, p in model_without_ddp.named_parameters()
if not match_name_keywords(n, args.lr_backbone_names) and \
not match_name_keywords(n, args.lr_linear_proj_names) and \
not match_name_keywords(n, args.embedding_related_names) and p.requires_grad],
"lr": args.lr,
"initial_lr": args.lr,
},
{
"params": [p for n, p in model_without_ddp.named_parameters() if match_name_keywords(n, args.lr_backbone_names) and p.requires_grad],
"lr": args.lr_backbone,
"initial_lr": args.lr_backbone,
},
]
optimizer = torch.optim.AdamW(param_dicts, weight_decay=args.weight_decay)
lr_scheduler = WarmupMultiStepLR(optimizer,
args.lr_drop_milestones,
gamma=0.1,
warmup_epochs=args.warmup_epochs,
warmup_factor=args.warmup_factor,
warmup_method='linear',
last_epoch=args.start_epoch - 1)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu], find_unused_parameters=True)
model_without_ddp = model.module
if args.dataset_file == "coco_panoptic":
# We also evaluate AP during panoptic training, on original coco DS
coco_val = datasets.dataset.build("val", args)
base_ds = get_coco_api_from_dataset(coco_val)
else:
base_ds = get_coco_api_from_dataset(dataset_val)
output_dir = Path(args.output_dir)
if args.resume:
if args.resume.startswith('https'):
checkpoint = torch.hub.load_state_dict_from_url(args.resume,
map_location='cpu',
check_hash=True)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
missing_keys, unexpected_keys = model_without_ddp.load_state_dict(
checkpoint['model'], strict=False)
unexpected_keys = [
k for k in unexpected_keys
if not (k.endswith('total_params') or k.endswith('total_ops'))
]
if len(missing_keys) > 0:
print('Missing Keys: {}'.format(missing_keys))
if len(unexpected_keys) > 0:
print('Unexpected Keys: {}'.format(unexpected_keys))
if args.fewshot_finetune:
if args.category_codes_cls_loss:
# Re-init weights of novel categories for few-shot finetune
novel_class_ids = datasets.get_class_ids(args.dataset_file,
type='novel')
if args.num_feature_levels == 1:
for novel_class_id in novel_class_ids:
nn.init.normal_(model_without_ddp.category_codes_cls.L.
weight[novel_class_id])
elif args.num_feature_levels > 1:
for classifier in model_without_ddp.category_codes_cls:
for novel_class_id in novel_class_ids:
nn.init.normal_(
classifier.L.weight[novel_class_id])
else:
raise RuntimeError
if args.eval:
# Evaluate only base categories
test_stats, coco_evaluator = evaluate(args,
model,
criterion,
postprocessors,
loader_val,
loader_support,
base_ds,
device,
type='base')
if args.output_dir:
utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval,
output_dir / "eval_base.pth")
# Evaluate only novel categories
test_stats, coco_evaluator = evaluate(args,
model,
criterion,
postprocessors,
loader_val,
loader_support,
base_ds,
device,
type='novel')
if args.output_dir:
utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval,
output_dir / "eval_novel.pth")
return
print("Start training...")
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
sampler_train.set_epoch(epoch)
train_stats = train_one_epoch(args, model, criterion, loader_train,
optimizer, device, epoch,
args.clip_max_norm)
lr_scheduler.step()
# Saving Checkpoints after each epoch
if args.output_dir and (not args.fewshot_finetune):
checkpoint_paths = [output_dir / 'checkpoint.pth']
for checkpoint_path in checkpoint_paths:
utils.save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'args': args,
}, checkpoint_path)
# Saving Checkpoints every args.save_every_epoch epoch(s)
if args.output_dir:
checkpoint_paths = []
if (epoch + 1) % args.save_every_epoch == 0:
checkpoint_paths.append(output_dir /
f'checkpoint{epoch:04}.pth')
for checkpoint_path in checkpoint_paths:
utils.save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'args': args,
}, checkpoint_path)
# Evaluation and Logging
if (epoch + 1) % args.eval_every_epoch == 0:
if 'base' in args.dataset_file:
evaltype = 'base'
else:
evaltype = 'all'
if args.fewshot_finetune:
evaltype = 'novel'
test_stats, coco_evaluator = evaluate(args,
model,
criterion,
postprocessors,
loader_val,
loader_support,
base_ds,
device,
type=evaltype)
log_stats = {
**{f'train_{k}': v
for k, v in train_stats.items()},
**{f'test_{k}': v
for k, v in test_stats.items()}, 'epoch': epoch,
'n_parameters': n_parameters,
'evaltype': evaltype
}
if args.output_dir and utils.is_main_process():
with (output_dir / "results.txt").open("a") as f:
f.write(json.dumps(log_stats) + "\n")
# for evaluation logs
if coco_evaluator is not None:
(output_dir / 'eval').mkdir(exist_ok=True)
if "bbox" in coco_evaluator.coco_eval:
filenames = ['latest.pth']
filenames.append(f'{epoch:03}.pth')
for name in filenames:
torch.save(coco_evaluator.coco_eval["bbox"].eval,
output_dir / "eval" / name)
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))