def __init__(self, args, train_backbone: bool, return_interm_layers: bool,
use_fpn: bool):
if args.model_arch.lower() == "vanilla":
backbone = getattr(torchvision.models, args.backbone)(
replace_stride_with_dilation=[False, False, args.dilation],
pretrained=is_main_process(),
norm_layer=FrozenBatchNorm2d)
elif args.model_arch.lower() in ["fpn", "fpn_v1"]:
backbone = resnet_fpn_backbone(args.backbone, is_main_process())
else:
raise NotImplementedError()
num_channels = 512 if args.backbone in ('resnet18',
'resnet34') else 2048
super().__init__(backbone, train_backbone, num_channels,
return_interm_layers, use_fpn)
def __init__(self, name: str, train_backbone: bool,
return_interm_layers: bool, dilation: bool,
training_spec: str, auto_checkpoint: str):
if name.startswith('autonet'):
# training_spec = 'sandbox/williamz/detr/res_autonet/autonet_training_spec.yaml'
# training_spec = os.path.join(os.environ["HOME"],'datasets/specs/autonet_training_spec.yaml')
training_spec = load_spec(training_spec)
model = deserialize_object(training_spec["model"])
# autonet checkpoint
# checkpoint = 'sandbox/williamz/detr/res_autonet/final_epoch.checkpoint'
checkpoint = auto_checkpoint
# checkpoint = os.path.join(os.environ["HOME"],'datasets/autonet/final_epoch.checkpoint')
if checkpoint is not None and os.path.isfile(
checkpoint) and is_main_process():
print(f'---------- Loading checkpoint for AutoNet -----')
loaded_states = torch.load(checkpoint)
model_state = loaded_states["model_state"]
model.load_state_dict(model_state, strict=False)
# backbone = model
else:
print(f'---------- No checkpoint for AutoNet -----')
# get drivenet
# IPython.embed()
modules = []
for block in model._blocks:
if 'drive2d' in block["task_name"]:
modules.append(getattr(model, block['name']))
backbone = nn.Sequential(*modules[:-1])
num_channels = 256
super().__init__(backbone, train_backbone, num_channels,
return_interm_layers)
def summarize(self):
if utils.is_main_process():
json_data = {"annotations": self.predictions}
predictions_json = os.path.join(self.output_dir, "predictions.json")
with open(predictions_json, "w") as f:
f.write(json.dumps(json_data))
return pq_compute(self.gt_json, predictions_json, gt_folder=self.gt_folder, pred_folder=self.output_dir)
return None
def __init__(self, name: str, train_backbone: bool,
return_interm_layers: bool, dilation: bool):
backbone = getattr(torchvision.models, name)(
replace_stride_with_dilation=[False, False, dilation],
pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d)
num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
super().__init__(backbone, train_backbone, num_channels,
return_interm_layers)
def __init__(self, ann_file, ann_folder, output_dir="panoptic_eval"):
self.gt_json = ann_file
self.gt_folder = ann_folder
if utils.is_main_process():
if not os.path.exists(output_dir):
os.mkdir(output_dir)
self.output_dir = output_dir
self.predictions = []
def __init__(self, name: str,
train_backbone: bool):
backbone = build_hrnet(name, pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d)
# [48, 96, 192, 384]
# [32, 64, 128, 256]
# [18, 36, 72, 144]
num_channels = 8 * int(name[-2:])
super().__init__(backbone, train_backbone, num_channels)
def __init__(self, name: str, train_backbone: bool,
return_interm_layers: bool, dilation: bool):
backbone = getattr(torchvision.models, name)(
replace_stride_with_dilation=[False, False, dilation],
pretrained=is_main_process(),
norm_layer=FrozenBatchNorm2d)
#backbone.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
super().__init__(backbone, train_backbone, num_channels,
return_interm_layers)
def display(self, batch):
entries = [self.prefix + self.batch_fmtstr.format(batch)]
entries += [str(meter) for meter in self.meters]
print('\t'.join(entries))
##################################
# Save to logging
##################################
if utils.is_main_process():
logging.info('\t'.join(entries))
def __init__(self, name: str, train_backbone: bool,
return_interm_layers: bool, dilation: bool):
norm_layer = FrozenBatchNorm2d
if name.startswith("resnet"):
backbone = getattr(torchvision.models, name)(
replace_stride_with_dilation=[False, False, dilation],
pretrained=is_main_process(),
norm_layer=norm_layer)
else:
assert name == "mobilenet_v2", f"Backbone {name} not supported"
backbone = getattr(torchvision.models,
name)(pretrained=is_main_process(),
norm_layer=norm_layer)
assert name not in ('resnet18',
'resnet34'), "number of channels are hard coded"
super().__init__(backbone, train_backbone, return_interm_layers)
if dilation:
self.strides[-1] = self.strides[-1] // 2
def __init__(
self,
name: str,
train_backbone: bool,
return_interm_layers: bool,
dilation: bool,
):
if name not in ["senet256"]:
optname = None
backbone = getattr(torchvision.models, name)(
replace_stride_with_dilation=[False, False, dilation],
pretrained=is_main_process(),
norm_layer=FrozenBatchNorm2d,
)
num_channels = 512 if name in ("resnet18", "resnet34") else 2048
else:
optname = "feat_extract"
backbone = senet50_256(pretrained=is_main_process())
num_channels = 2048
super().__init__(backbone, train_backbone, num_channels,
return_interm_layers, optname)
def __init__(self, name: str,
train_backbone: bool,
return_interm_layers: bool,
dilation: bool):
# get the function with name of backbone model, i.e. call the constructor,
# and replace some constructor inputs, i.e. dilation flag, whether to use pretrained weights,
# replace normalization layers within to fix nan bug
backbone = getattr(torchvision.models, name)(
replace_stride_with_dilation=[False, False, dilation],
pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d)
num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
# call parent init to return interim layers for interim mask output
super().__init__(backbone, train_backbone, num_channels, return_interm_layers)
def __init__(self, name: str, train_backbone: bool,
return_interm_layers: bool, dilation: bool):
backbone = getattr(torchvision.models, name)(
replace_stride_with_dilation=[False, False, dilation],
pretrained=is_main_process(),
norm_layer=FrozenBatchNorm2d)
print("\n === Convert All Conv/MaxPool layers to Sphere Layers ===")
self.converLayers(backbone)
print("done!")
num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
super().__init__(backbone, train_backbone, num_channels,
return_interm_layers)
def __init__(self, name: str,
train_backbone: bool,
return_interm_layers: bool,
dilation: bool):
# if name == 'EfficientNetB0':
# backbone = EfficientNet.from_pretrained('efficientnet-b0')
# separateblocksbo(backbone)
# num_channels = 320
# else:
backbone = getattr(torchvision.models, name)(
replace_stride_with_dilation= [False, False, dilation],
pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d)
num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
super().__init__(backbone, train_backbone, num_channels, return_interm_layers)
def __init__(self, name: str, train_backbone: bool,
return_interm_layers: bool, args):
self.args = args
dilation = args.dilation
norm_layer = FrozenBatchNorm2d
backbone = getattr(torchvision.models, name)(
replace_stride_with_dilation=[False, False, dilation],
pretrained=is_main_process(),
norm_layer=norm_layer)
assert name not in (
'resnet18', 'resnet34'
), "number of channels are hard coded, cannot use res18 & res34."
super().__init__(backbone, train_backbone, return_interm_layers, args)
if dilation:
self.strides[-1] = self.strides[-1] // 2
def __init__(self, name: str, train_backbone: bool, return_layers: List,
dilation: bool):
if dilation:
dilation = [False, True, True] # workaround to achieve stride of 8
else:
dilation = [False, False, False]
backbone = getattr(torchvision.models,
name)(replace_stride_with_dilation=dilation,
pretrained=is_main_process(),
norm_layer=FrozenBatchNorm2d)
super().__init__(name, backbone, train_backbone, return_layers)
final_layer = int(return_layers[-1][-1])
self.dilation = dilation
self.stride = 4
for layer in range(final_layer - 1):
if not dilation[layer]:
self.stride = self.stride * 2
def validate(val_loader,
model,
criterion,
args,
FAST_VALIDATE_FLAG=False,
DUMP_RESULT=False):
batch_time = AverageMeter('Time', ':6.3f')
program_acc = AverageMeter('Acc@Program', ':6.2f')
program_group_acc = AverageMeter('Acc@ProgramGroup', ':4.2f')
program_non_empty_acc = AverageMeter('Acc@ProgramNonEmpty', ':4.2f')
# bitmap_precision = AverageMeter('Precision@Bitmap', ':4.2f')
# bitmap_recall = AverageMeter('Recall@Bitmap', ':4.2f')
# full_answer_acc = AverageMeter('Acc@Full', ':6.2f')
short_answer_acc = AverageMeter('Acc@Short', ':6.2f')
progress = ProgressMeter(len(val_loader), [
batch_time, program_acc, program_group_acc, program_non_empty_acc,
short_answer_acc
],
prefix='Test: ')
# switch to evaluate mode
model.eval()
if DUMP_RESULT:
quesid2ans = {}
with torch.no_grad():
end = time.time()
for i, (data_batch) in enumerate(val_loader):
questionID, questions, gt_scene_graphs, programs, full_answers, short_answer_label, types = data_batch
questions, gt_scene_graphs, programs, full_answers, short_answer_label = [
datum.to(device=cuda, non_blocking=True) for datum in [
questions, gt_scene_graphs, programs, full_answers,
short_answer_label
]
]
this_batch_size = questions.size(1)
if FAST_VALIDATE_FLAG:
raise NotImplementedError(
"Should not use fast validation. Only for short answer accuracy"
)
##################################
# Prepare training input and training target for text generation
##################################
programs_input = programs[:-1]
programs_target = programs[1:]
full_answers_input = full_answers[:-1]
full_answers_target = full_answers[1:]
##################################
# Forward evaluate data
##################################
output = model(questions, gt_scene_graphs, programs_input,
full_answers_input)
programs_output, short_answer_logits = output
##################################
# Convert output probability to top1 guess
# So that we could measure accuracy
##################################
programs_output_pred = programs_output.detach().topk(
k=1, dim=-1, largest=True, sorted=True)[1].squeeze(-1)
# full_answers_output_pred = full_answers_output.detach().topk(
# k=1, dim=-1, largest=True, sorted=True
# )[1].squeeze(-1)
else:
programs_target = programs
full_answers_target = full_answers
##################################
# Greedy decoding-based evaluation
##################################
output = model(questions,
gt_scene_graphs,
None,
None,
SAMPLE_FLAG=True)
programs_output_pred, short_answer_logits = output
##################################
# Neural Execution Engine Bitmap loss
# ground truth stored at gt_scene_graphs.y
# using torch.nn.BCELoss - torch.nn.functional.binary_cross_entropy
##################################
# precision, precision_div, recall, recall_div = bitmap_precision_recall(
# execution_bitmap, gt_scene_graphs.y, threshold=0.5
# )
# bitmap_precision.update(precision, precision_div)
# bitmap_recall.update(recall, recall_div)
##################################
# Calculate Fast Evaluation for each module
##################################
this_short_answer_acc1 = accuracy(short_answer_logits.detach(),
short_answer_label,
topk=(1, ))
short_answer_acc.update(this_short_answer_acc1[0].item(),
this_batch_size)
text_pad_idx = GQATorchDataset.TEXT.vocab.stoi[
GQATorchDataset.TEXT.pad_token]
this_program_acc, this_program_group_acc, this_program_non_empty_acc = program_string_exact_match_acc(
programs_output_pred,
programs_target,
padding_idx=text_pad_idx,
group_accuracy_WAY_NUM=GQATorchDataset.MAX_EXECUTION_STEP)
program_acc.update(this_program_acc, this_batch_size)
program_group_acc.update(
this_program_group_acc,
this_batch_size // GQATorchDataset.MAX_EXECUTION_STEP)
program_non_empty_acc.update(this_program_non_empty_acc,
this_batch_size)
# this_full_answers_acc = string_exact_match_acc(
# full_answers_output_pred.detach(), full_answers_target, padding_idx=text_pad_idx
# )
# full_answer_acc.update(this_full_answers_acc, this_batch_size)
##################################
# Example Visualization from the first batch
##################################
if i == 0 and True:
for batch_idx in range(min(this_batch_size, 128)):
##################################
# print Question and Question ID
##################################
question = questions[:, batch_idx].cpu()
question_sent, _ = GQATorchDataset.indices_to_string(
question, True)
print(
"Question({}) QID({}):".format(batch_idx,
questionID[batch_idx]),
question_sent)
if utils.is_main_process():
logging.info("Question({}) QID({}): {}".format(
batch_idx, questionID[batch_idx], question_sent))
##################################
# print program prediction
##################################
for instr_idx in range(GQATorchDataset.MAX_EXECUTION_STEP):
true_batch_idx = instr_idx + GQATorchDataset.MAX_EXECUTION_STEP * batch_idx
gt = programs[:, true_batch_idx].cpu()
pred = programs_output_pred[:, true_batch_idx]
pred_sent, _ = GQATorchDataset.indices_to_string(
pred, True)
gt_sent, _ = GQATorchDataset.indices_to_string(
gt, True)
if len(pred_sent) == 0 and len(gt_sent) == 0:
# skip if both target and prediciton are empty
continue
# gt_caption
print(
"Generated Program ({}): ".format(true_batch_idx),
pred_sent, " Ground Truth Program ({}):".format(
true_batch_idx), gt_sent)
if utils.is_main_process():
# gt_caption
logging.info(
"Generated Program ({}): {} Ground Truth Program ({}): {}"
.format(true_batch_idx, pred_sent,
true_batch_idx, gt_sent))
##################################
# print full answer prediction
##################################
# gt = full_answers[:, batch_idx].cpu()
# pred = full_answers_output_pred[:, batch_idx]
# pred_sent, _ = GQATorchDataset.indices_to_string(pred, True)
# gt_sent, _ = GQATorchDataset.indices_to_string(gt, True)
# # gt_caption
# print(
# "Generated Full Answer ({}): ".format(batch_idx), pred_sent,
# "Ground Truth Full Answer ({}):".format(batch_idx), gt_sent
# )
# if utils.is_main_process():
# # gt_caption
# logging.info("Generated Full Answer ({}): {} Ground Truth Full Answer ({}): {}".format(
# batch_idx, pred_sent, batch_idx, gt_sent
# ))
##################################
# Dump Results if enabled
##################################
if DUMP_RESULT:
short_answer_pred_score, short_answer_pred_label = short_answer_logits.max(
1)
short_answer_pred_score, short_answer_pred_label = short_answer_pred_score.cpu(
), short_answer_pred_label.cpu()
for batch_idx in range(this_batch_size):
##################################
# print Question and Question ID
##################################
question = questions[:, batch_idx].cpu()
question_sent, _ = GQATorchDataset.indices_to_string(
question, True)
##################################
# print program prediction
##################################
ground_truth_program_list = []
predicted_program_list = []
for instr_idx in range(GQATorchDataset.MAX_EXECUTION_STEP):
true_batch_idx = instr_idx + GQATorchDataset.MAX_EXECUTION_STEP * batch_idx
gt = programs[:, true_batch_idx].cpu()
pred = programs_output_pred[:, true_batch_idx]
pred_sent, _ = GQATorchDataset.indices_to_string(
pred, True)
gt_sent, _ = GQATorchDataset.indices_to_string(
gt, True)
if len(pred_sent) == 0 and len(gt_sent) == 0:
# skip if both target and prediciton are empty
continue
ground_truth_program_list.append(gt_sent)
predicted_program_list.append(pred_sent)
##################################
# print full answer prediction
##################################
# gt = full_answers[:, batch_idx].cpu()
# pred = full_answers_output_pred[:, batch_idx]
# pred_sent, _ = GQATorchDataset.indices_to_string(pred, True)
# gt_sent, _ = GQATorchDataset.indices_to_string(gt, True)
# gt_caption
##################################
# get short answer prediction
##################################
qid = questionID[batch_idx]
quesid2ans[qid] = {
"questionId":
str(qid),
"question":
question_sent,
"ground_truth_program_list":
ground_truth_program_list,
"predicted_program_list":
predicted_program_list,
"answer":
GQATorchDataset.label2ans[
short_answer_label[batch_idx].cpu().item()],
# predicted short answer
"prediction":
GQATorchDataset.label2ans[
short_answer_pred_label[batch_idx].cpu().item()],
"prediction_score":
'{:.2f}'.format(
short_answer_pred_score[batch_idx].cpu().item()),
"types":
types[batch_idx],
}
##################################
# measure elapsed time
##################################
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0 or i == len(val_loader) - 1:
progress.display(i)
##################################
# Only for dubugging: short cut the evaluation loop
##################################
# break
##################################
# Give final score
##################################
progress.display(batch=len(val_loader))
if DUMP_RESULT:
result_dump_path = os.path.join(args.output_dir, "dump_results.json")
with open(result_dump_path, 'w') as f:
json.dump(quesid2ans, f, indent=4, sort_keys=True)
print("Result Dumped!", str(result_dump_path))
return
def main(args):
# args = parser.parse_args()
utils.init_distributed_mode(args)
print("git:\n {}\n".format(utils.get_sha()))
print(args)
if args.seed is not None:
# random.seed(args.seed)
# torch.manual_seed(args.seed)
cudnn.deterministic = True
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
# fix the seed for reproducibility
seed = args.seed + utils.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
##################################
# Logging setting
##################################
if args.output_dir and utils.is_main_process():
logging.basicConfig(
filename=os.path.join(args.output_dir, args.log_name),
filemode='w',
format=
'%(asctime)s: %(levelname)s: [%(filename)s:%(lineno)d]: %(message)s',
level=logging.INFO)
warnings.filterwarnings("ignore")
##################################
# Save to logging
##################################
if utils.is_main_process():
logging.info(str(args))
##################################
# Initialize dataset
##################################
if not args.evaluate:
# build_vocab_flag=True, # Takes a long time to build a vocab
train_dataset = GQATorchDataset(split='train_unbiased',
build_vocab_flag=False,
load_vocab_flag=False)
if args.distributed:
sampler_train = torch.utils.data.DistributedSampler(train_dataset)
else:
sampler_train = torch.utils.data.RandomSampler(train_dataset)
batch_sampler_train = torch.utils.data.BatchSampler(sampler_train,
args.batch_size,
drop_last=True)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_sampler=batch_sampler_train,
collate_fn=GQATorchDataset_collate_fn,
num_workers=args.workers)
# Old version
# train_loader = torch.utils.data.DataLoader(
# train_dataset, batch_size=args.batch_size, shuffle=True,
# collate_fn=GQATorchDataset_collate_fn,
# num_workers=args.workers, pin_memory=True)
val_dataset_list = []
for eval_split in args.evaluate_sets:
val_dataset_list.append(
GQATorchDataset(split=eval_split,
build_vocab_flag=False,
load_vocab_flag=args.evaluate))
val_dataset = torch.utils.data.ConcatDataset(val_dataset_list)
if args.distributed:
sampler_val = torch.utils.data.DistributedSampler(val_dataset,
shuffle=False)
else:
sampler_val = torch.utils.data.SequentialSampler(val_dataset)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=args.batch_size,
sampler=sampler_val,
drop_last=False,
collate_fn=GQATorchDataset_collate_fn,
num_workers=args.workers)
# Old version
# val_loader = torch.utils.data.DataLoader(
# val_dataset,
# batch_size=args.batch_size, shuffle=False,
# collate_fn=GQATorchDataset_collate_fn,
# num_workers=args.workers, pin_memory=True)
##################################
# Initialize model
# - note: must init dataset first. Since we will use the vocab from the dataset
##################################
model = PipelineModel()
##################################
# Deploy model on GPU
##################################
model = model.to(device=cuda)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu], find_unused_parameters=True)
model_without_ddp = model.module
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
##################################
# define optimizer (and scheduler)
##################################
# optimizer = torch.optim.SGD(model.parameters(), args.lr,
# momentum=args.momentum,
# weight_decay=args.weight_decay)
optimizer = torch.optim.Adam(
params=model.parameters(),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0, # weight_decay=args.weight_decay
amsgrad=False,
)
# optimizer = torch.optim.AdamW(
# params=model.parameters(),
# lr=args.lr,
# weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
model_without_ddp.load_state_dict(checkpoint['model'])
if not args.evaluate:
if 'optimizer' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
if 'lr_scheduler' in checkpoint:
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
if 'epoch' in checkpoint:
args.start_epoch = checkpoint['epoch'] + 1
# checkpoint = torch.load(args.resume)
# args.start_epoch = checkpoint['epoch']
# model.load_state_dict(checkpoint['state_dict'])
# optimizer.load_state_dict(checkpoint['optimizer'])
# print("=> loaded checkpoint '{}' (epoch {})"
# .format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# cudnn.benchmark = True
##################################
# Define loss functions (criterion)
##################################
# criterion = torch.nn.CrossEntropyLoss().cuda()
text_pad_idx = GQATorchDataset.TEXT.vocab.stoi[
GQATorchDataset.TEXT.pad_token]
criterion = {
"program":
torch.nn.CrossEntropyLoss(ignore_index=text_pad_idx).to(device=cuda),
"full_answer":
torch.nn.CrossEntropyLoss(ignore_index=text_pad_idx).to(device=cuda),
"short_answer":
torch.nn.CrossEntropyLoss().to(device=cuda),
# "short_answer": torch.nn.BCEWithLogitsLoss().to(device=cuda), # sigmoid
"execution_bitmap":
torch.nn.BCELoss().to(device=cuda),
}
##################################
# If Evaluate Only
##################################
if args.evaluate:
validate(val_loader, model, criterion, args, DUMP_RESULT=True)
return
##################################
# Main Training Loop
##################################
# best_acc1 = 0
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
##################################
# In distributed mode, calling the :meth`set_epoch(epoch) <set_epoch>` method
# at the beginning of each epoch before creating the DataLoader iterator is necessary
# to make shuffling work properly across multiple epochs.
# Otherwise, the same ordering will be always used.
##################################
sampler_train.set_epoch(epoch)
lr_scheduler.step()
# adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, args)
# evaluate on validation set
if (epoch + 1) % 5 == 0:
validate(val_loader,
model,
criterion,
args,
FAST_VALIDATE_FLAG=False)
# # remember best acc@1 and save checkpoint
# save_checkpoint({
# 'epoch': epoch + 1,
# # 'arch': args.arch,
# 'state_dict': model.state_dict(),
# # 'best_acc1': best_acc1,
# 'optimizer' : optimizer.state_dict(),
# }, is_best)
if args.output_dir:
output_dir = pathlib.Path(args.output_dir)
checkpoint_paths = [output_dir / 'checkpoint.pth']
# extra checkpoint before LR drop and every 100 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 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)
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
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
param_dicts = [
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" not in n and p.requires_grad
]
},
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" in n and p.requires_grad
],
"lr":
args.lr_backbone,
},
]
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)
dataset_train = build_dataset(image_set='train', args=args)
dataset_val = build_dataset(image_set='val', args=args)
if args.distributed:
sampler_train = DistributedSampler(dataset_train)
sampler_val = 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)
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)
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')
model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
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
#cab
writer = SummaryWriter("runs/" + args.tb_name)
best_value = 0
print("Start training, best_value is " + str(best_value))
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()
test_stats, coco_evaluator = evaluate(model, criterion, postprocessors,
data_loader_val, base_ds, device,
args.output_dir)
#cab
for k, v in train_stats.items():
if isinstance(v, float):
writer.add_scalar(f'train_{k}', v, epoch)
new_value = 0
for k, v in test_stats.items():
if (isinstance(v, float)):
writer.add_scalar(f'test_{k}', v, epoch)
if (k == "coco_eval_bbox"):
new_value = v[0]
writer.add_scalar(
'Bbox Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ]',
v[0], epoch)
writer.add_scalar(
'Bbox Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ]',
v[1], epoch)
writer.add_scalar(
'Bbox Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ]',
v[2], epoch)
writer.add_scalar(
'Bbox Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ]',
v[3], epoch)
writer.add_scalar(
'Bbox Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ]',
v[4], epoch)
writer.add_scalar(
'Bbox Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ]',
v[5], epoch)
writer.add_scalar(
'Bbox Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ]',
v[6], epoch)
writer.add_scalar(
'Bbox Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ]',
v[7], epoch)
writer.add_scalar(
'Bbox Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ]',
v[8], epoch)
writer.add_scalar(
'Bbox Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ]',
v[9], epoch)
writer.add_scalar(
'Bbox Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ]',
v[10], epoch)
writer.add_scalar(
'Bbox Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ]',
v[11], epoch)
if (k == "coco_eval_masks"):
new_value = v[0]
writer.add_scalar(
'Mask Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ]',
v[0], epoch)
writer.add_scalar(
'Mask Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ]',
v[1], epoch)
writer.add_scalar(
'Mask Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ]',
v[2], epoch)
writer.add_scalar(
'Mask Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ]',
v[3], epoch)
writer.add_scalar(
'Mask Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ]',
v[4], epoch)
writer.add_scalar(
'Mask Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ]',
v[5], epoch)
writer.add_scalar(
'Mask Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ]',
v[6], epoch)
writer.add_scalar(
'Mask Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ]',
v[7], epoch)
writer.add_scalar(
'Mask Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ]',
v[8], epoch)
writer.add_scalar(
'Mask Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ]',
v[9], epoch)
writer.add_scalar(
'Mask Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ]',
v[10], epoch)
writer.add_scalar(
'Mask Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ]',
v[11], epoch)
print("Epoch finished, best_value is " + str(best_value))
save_pth = False
if best_value < new_value:
best_value = new_value
save_pth = True
if args.output_dir:
checkpoint_paths = [output_dir / 'checkpoint.pth']
# extra checkpoint before LR drop and every 100 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 0:
checkpoint_paths.append(output_dir /
f'checkpoint{epoch:04}.pth')
if save_pth:
checkpoint_paths.append(output_dir / f'best.pth')
bestLog = open(output_dir / 'best_log.txt', 'w+')
bestLog.write(f'Saved model at epoch {epoch:04}\n')
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)
#/cab
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()))
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 = build_model(args)
model.to(device)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
model_without_ddp = model.module
n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
print('number of params:', n_parameters)
param_dicts = [
{"params": [p for n, p in model_without_ddp.named_parameters() if "backbone" not in n and p.requires_grad]},
{
"params": [p for n, p in model_without_ddp.named_parameters() if "backbone" in n and p.requires_grad],
"lr": args.lr_backbone,
},
]
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)
dataset_train = build_dataset(image_set='train', args=args)
if args.distributed:
sampler_train = DistributedSampler(dataset_train)
else:
sampler_train = torch.utils.data.RandomSampler(dataset_train)
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)
# Load from pretrained DETR model.
assert args.num_queries == 100, args.num_queries
assert args.enc_layers == 6 and args.dec_layers == 6
assert args.backbone in ['resnet50', 'resnet101', 'swin'], args.backbone
if args.backbone == 'resnet50':
pretrain_model = './data/detr_coco/detr-r50-e632da11.pth'
elif args.backbone == 'resnet101':
pretrain_model = './data/detr_coco/detr-r101-2c7b67e5.pth'
else:
pretrain_model = None
if pretrain_model is not None:
pretrain_dict = torch.load(pretrain_model, map_location='cpu')['model']
my_model_dict = model_without_ddp.state_dict()
pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in my_model_dict}
my_model_dict.update(pretrain_dict)
model_without_ddp.load_state_dict(my_model_dict)
output_dir = Path(args.output_dir)
if args.resume:
checkpoint = torch.load(args.resume, map_location='cpu')
model_without_ddp.load_state_dict(checkpoint['model'])
if 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
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 10 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 0:
checkpoint_paths.append(output_dir / f'checkpoint{epoch:04}.pth')
if (epoch + 1) > args.lr_drop and (epoch + 1) % 10 == 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 args.output_dir and utils.is_main_process():
with (output_dir / "log.txt").open("a") as f:
f.write(json.dumps(log_stats) + "\n")
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)
# IPython.embed()
# IPython.embed()
# os.system("sudo chmod -R 777 /home/shuxuang/.cache/")
model, criterion, postprocessors = build_model(
args) # use the same model as detr paper on coco
model.to(device)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
param_dicts = [
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" not in n and p.requires_grad
]
},
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" in n and p.requires_grad
],
"lr":
args.lr_backbone,
},
]
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)
# dataset_train = build_dataset(image_set='train', args=args)
# dataset_val = build_dataset(image_set='val', args=args)
# modify the dataset from coco to nvdata
# home_dir = os.environ["HOME"]
# dataset_train_ = build_nvdataset(dataset_root=[
# os.path.join(os.environ["HOME"],'datasets/annotation_sql_nvidia'),
# os.path.join(os.environ["HOME"], 'datasets/frames_nvidia')],
# mode='train')
# dataset_val = build_nvdataset(dataset_root=[
# os.path.join(os.environ["HOME"],'datasets/test'),
# os.path.join(os.environ["HOME"], 'datasets/frames_nvidia')],
# mode='test')
# indices_50k =np.load(os.path.join(os.environ["HOME"],'datasets/id_1_criterion_Max_SSD_num_labels_50000.npy'))
dataset_train = build_nvdataset(
dataset_root=[args.dataset_root_sql, args.dataset_root_img],
mode='train',
camera=args.camera)
dataset_val = build_nvdataset(
dataset_root=[args.dataset_root_test, args.dataset_root_test],
mode='test',
camera=args.camera)
if args.root_indices is not None:
indices_50k = np.load(os.path.join(args.root_indices))
# indices_50k =np.load(os.path.join(os.environ["HOME"],'datasets/id_1_criterion_Max_SSD_num_labels_50000.npy'))
dataset_train = Subset(dataset_train, indices_50k)
# IPython.embed()
print("Train samples: %d" % (len(dataset_train)))
if args.distributed:
sampler_train = DistributedSampler(dataset_train)
sampler_val = 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)
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)
# 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)
# elif args.dataset_file == "nvdata":
# 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')
model_without_ddp.load_state_dict(checkpoint['model'])
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
# if args.eval:
# test_stats, coco_evaluator = evaluate_nvdata(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.eval:
# evaluate(model, dataset_val, postprocessors, device)
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 100 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 50 == 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_nvdata(
# 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}
log_stats = {
**{f'train_{k}': v
for k, v in train_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 evaluate(model,
criterion,
postprocessors,
data_loader,
base_ds,
device,
output_dir,
log_step=0):
model.eval()
criterion.eval()
metric_logger = utils.MetricLogger(delimiter=" ")
metric_logger.add_meter(
'class_error', utils.SmoothedValue(window_size=1, fmt='{value:.2f}'))
header = 'Test:'
iou_types = tuple(k for k in ('segm', 'bbox')
if k in postprocessors.keys())
coco_evaluator = CocoEvaluator(base_ds, iou_types)
panoptic_evaluator = None
if 'panoptic' in postprocessors.keys():
panoptic_evaluator = PanopticEvaluator(
data_loader.dataset.ann_file,
data_loader.dataset.ann_folder,
output_dir=os.path.join(output_dir, "panoptic_eval"),
)
dataset = data_loader.dataset
classes = {
cat["id"]: cat["name"]
for cat in dataset.coco.dataset["categories"]
}
wandb_imgs = {"images": [], "self_attention": [], "attention": []}
# Log every 50 steps and in step 0
log_this = output_dir and utils.is_main_process() and (
(log_step + 1) % 50 == 0 or log_step == 0)
conv_features, enc_attn_weights, dec_attn_weights = [], [], []
for samples, targets in metric_logger.log_every(data_loader, 10, header):
samples = samples.to(device)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
log_image = False
if log_this:
if len(LOG_IDX) == 15:
if targets[0]["image_id"] in LOG_IDX:
log_image = True
elif random.random() < 0.3 and len(
targets[0]["labels"].tolist()) > 3:
LOG_IDX.append(targets[0]["image_id"])
log_image = True
if log_image:
# Taken from https://colab.research.google.com/github/facebookresearch/detr/blob/colab/notebooks/detr_attention.ipynb
hooks = [
model.module.backbone[-2].register_forward_hook(
lambda self, input, output: conv_features.append(output
)),
model.module.transformer.encoder.layers[-1].self_attn.
register_forward_hook(lambda self, input, output:
enc_attn_weights.append(output[1])),
model.module.transformer.decoder.layers[-1].multihead_attn.
register_forward_hook(lambda self, input, output:
dec_attn_weights.append(output[1])),
]
outputs = model(samples)
loss_dict = criterion(outputs, targets)
weight_dict = criterion.weight_dict
# reduce losses over all GPUs for logging purposes
loss_dict_reduced = utils.reduce_dict(loss_dict)
loss_dict_reduced_scaled = {
k: v * weight_dict[k]
for k, v in loss_dict_reduced.items() if k in weight_dict
}
loss_dict_reduced_unscaled = {
f'{k}_unscaled': v
for k, v in loss_dict_reduced.items()
}
metric_logger.update(loss=sum(loss_dict_reduced_scaled.values()),
**loss_dict_reduced_scaled,
**loss_dict_reduced_unscaled)
metric_logger.update(class_error=loss_dict_reduced['class_error'])
orig_target_sizes = torch.stack([t["orig_size"] for t in targets],
dim=0)
results = postprocessors['bbox'](outputs, orig_target_sizes)
# Gather images to log to wandb
if log_image:
# get the HxW shape of the feature maps of the CNN
f_map = conv_features[-1]['0'].tensors.cpu()
shape = f_map.shape[-2:]
sattn = enc_attn_weights[-1][0].reshape(shape + shape).cpu()
dec_att = dec_attn_weights[-1].cpu()
target = targets[0]
logits = outputs["pred_logits"][0]
boxes = outputs["pred_boxes"][0]
pred = {"pred_logits": logits, "pred_boxes": boxes}
name = dataset.coco.imgs[target["image_id"].item()]["file_name"]
path = os.path.join(dataset.root, name)
img, self_attention, att_map = create_wandb_img(
classes, path, target, pred, sattn, f_map, dec_att)
wandb_imgs["images"].append(img)
wandb_imgs["self_attention"].append(self_attention)
wandb_imgs["attention"].append(att_map)
# Free memory
del conv_features[-1]
del enc_attn_weights[-1]
for hook in hooks:
hook.remove()
if 'segm' in postprocessors.keys():
target_sizes = torch.stack([t["size"] for t in targets], dim=0)
results = postprocessors['segm'](results, outputs,
orig_target_sizes, target_sizes)
res = {
target['image_id'].item(): output
for target, output in zip(targets, results)
}
if coco_evaluator is not None:
coco_evaluator.update(res)
if panoptic_evaluator is not None:
res_pano = postprocessors["panoptic"](outputs, target_sizes,
orig_target_sizes)
for i, target in enumerate(targets):
image_id = target["image_id"].item()
file_name = f"{image_id:012d}.png"
res_pano[i]["image_id"] = image_id
res_pano[i]["file_name"] = file_name
panoptic_evaluator.update(res_pano)
# Log all images to wandb
if log_this:
wandb.log({"Images": wandb_imgs["images"]}, step=log_step)
wandb.log({"Self Attention": wandb_imgs["self_attention"]},
step=log_step)
wandb.log({"Attention": wandb_imgs["attention"]}, step=log_step)
# gather the stats from all processes
metric_logger.synchronize_between_processes()
print("Averaged stats:", metric_logger)
if coco_evaluator is not None:
coco_evaluator.synchronize_between_processes()
if panoptic_evaluator is not None:
panoptic_evaluator.synchronize_between_processes()
# accumulate predictions from all images
if coco_evaluator is not None:
coco_evaluator.accumulate()
coco_evaluator.summarize()
panoptic_res = None
if panoptic_evaluator is not None:
panoptic_res = panoptic_evaluator.summarize()
stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()}
if coco_evaluator is not None:
if 'bbox' in postprocessors.keys():
stats['coco_eval_bbox'] = coco_evaluator.coco_eval[
'bbox'].stats.tolist()
if 'segm' in postprocessors.keys():
stats['coco_eval_masks'] = coco_evaluator.coco_eval[
'segm'].stats.tolist()
if panoptic_res is not None:
stats['PQ_all'] = panoptic_res["All"]
stats['PQ_th'] = panoptic_res["Things"]
stats['PQ_st'] = panoptic_res["Stuff"]
return stats, coco_evaluator
def main(args):
utils.init_distributed_mode(args)
print("git:\n {}\n".format(utils.get_sha()))
wandb.init(project="qpic-project",
entity="sangbaeklee",
group="experiment_qpic")
wandb.config = {
"learning_rate": args.lr,
"epochs": args.epochs,
"batch_size": args.batch_size,
}
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)
wandb.watch(model)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
param_dicts = [
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" not in n and p.requires_grad
]
},
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" in n and p.requires_grad
],
"lr":
args.lr_backbone,
},
]
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)
dataset_train = build_dataset(image_set='train', args=args)
dataset_val = build_dataset(image_set='val', args=args)
if args.distributed:
sampler_train = DistributedSampler(dataset_train)
sampler_val = 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)
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)
if not args.hoi:
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')
model_without_ddp.load_state_dict(checkpoint['model'])
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
elif args.pretrained:
checkpoint = torch.load(args.pretrained, map_location='cpu')
model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
if args.eval:
if args.hoi:
test_stats = evaluate_hoi(args.dataset_file, model, postprocessors,
data_loader_val,
args.subject_category_id, device)
return
else:
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_paths = [output_dir / 'checkpoint.pth']
# extra checkpoint before LR drop and every 100 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 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)
if args.hoi:
test_stats = evaluate_hoi(args.dataset_file, model, postprocessors,
data_loader_val,
args.subject_category_id, device)
coco_evaluator = None
else:
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
}
#import pdb; pdb.set_trace()
if args.dataset_file == 'hico':
wandb.log({
"loss": train_stats['loss'],
"mAP": test_stats['mAP'],
"mAP rare": test_stats['mAP rare'],
"mAP non-rare": test_stats['mAP non-rare'],
"mean max recall": test_stats['mean max recall']
})
elif args.dataset_file == 'vcoco':
wandb.log({
"mAP_all": test_stats['mAP_all'],
"mAP_thesis": test_stats['mAP_thesis'],
"AP_hold_obj": test_stats['AP_hold_obj'],
"AP_stand": test_stats['AP_stand'],
"AP_sit_instr": test_stats['AP_sit_instr'],
"AP_ride_instr": test_stats['AP_ride_instr'],
"AP_walk": test_stats['AP_walk'],
"AP_look_obj": test_stats['AP_look_obj'],
"AP_hit_instr": test_stats['AP_hit_instr'],
"AP_hit_obj": test_stats['AP_hit_obj'],
"AP_eat_obj": test_stats['AP_eat_obj'],
"AP_eat_instr": test_stats['AP_eat_instr'],
"AP_jump_instr": test_stats['AP_jump_instr'],
"AP_lay_instr": test_stats['AP_lay_instr'],
"AP_talk_on_phone_instr": test_stats['AP_talk_on_phone_instr'],
"AP_carry_obj": test_stats['AP_carry_obj'],
"AP_throw_obj": test_stats['AP_throw_obj'],
"AP_catch_obj": test_stats['AP_catch_obj'],
"AP_cut_instr": test_stats['AP_cut_instr'],
"AP_cut_obj": test_stats['AP_cut_obj'],
"AP_run": test_stats['AP_run'],
"AP_work_on_computer_instr": test_stats['AP_work_on_computer_instr'],
"AP_ski_instr": test_stats['AP_ski_instr'],
"AP_surf_instr": test_stats['AP_surf_instr'],
"AP_skateboard_instr": test_stats['AP_skateboard_instr'],
"AP_smile": test_stats['AP_smile'],
"AP_drink_instr": test_stats['AP_drink_instr'],
"AP_kick_obj": test_stats['AP_kick_obj'],
"AP_point_instr": test_stats['AP_point_instr'],
"AP_read_obj": test_stats['AP_read_obj'],
"AP_snowboard_instr": test_stats['AP_snowboard_instr'],\
"loss" : train_stats['loss']
})
else:
continue
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)
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
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
param_dicts = [
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" not in n and p.requires_grad
]
},
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" in n and p.requires_grad
],
"lr":
args.lr_backbone,
},
]
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)
dataset_train = build_dataset(image_set='train', args=args)
dataset_val = build_dataset(image_set='val', args=args)
if args.distributed:
sampler_train = DistributedSampler(dataset_train)
sampler_val = 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)
data_loader_val = DataLoader(dataset_val,
batch_size=1,
sampler=sampler_val,
drop_last=False,
collate_fn=utils.collate_fn,
num_workers=args.num_workers)
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:
io.load_frozen(args, model_without_ddp)
output_dir = Path(args.output_dir)
if args.resume:
io.resume(args, model_without_ddp, optimizer, lr_scheduler)
elif args.finetune:
io.finetune(args, model_without_ddp)
if args.eval:
if args.output_dir and utils.is_main_process():
io.init_wandb(args.dataset_file + "-detr-eval",
model,
args,
n_parameters=n_parameters)
test_stats, evaluator = evaluate(model, criterion, postprocessors,
data_loader_val, base_ds, device,
args.output_dir)
if args.output_dir:
io.save_on_master(evaluator.coco_eval["bbox"].eval,
output_dir / "eval.pth")
return
print("Start training")
start_time = time.time()
if args.output_dir and utils.is_main_process():
io.init_wandb(args.dataset_file + "-detr",
model,
args,
n_parameters=n_parameters)
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:
# extra checkpoint before LR drop and every 100 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 0:
io.save_checkpoint(args, model_without_ddp, optimizer,
lr_scheduler, epoch)
test_stats, evaluator = evaluate(model, criterion, postprocessors,
data_loader_val, base_ds, device,
args.output_dir, epoch)
if utils.is_main_process() and args.output_dir:
io.log_wandb(train_stats, test_stats)
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
# save final model
if utils.is_main_process() and args.output_dir:
io.save_on_master(model_without_ddp, output_dir / "model_final.pth")
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)
# 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()))
# align with DETR format
args.dataset_file = 'ImageNet'
args.masks = None
# freeze cnn weights
args.lr_backbone = 0 if args.fre_cnn else args.lr
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
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
param_dicts = [
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" not in n and p.requires_grad
]
},
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" in n and p.requires_grad
],
"lr":
args.lr_backbone,
},
]
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)
dataset_train = build_dataset(image_set='train', args=args)
if args.distributed:
sampler_train = DistributedSampler(dataset_train)
else:
sampler_train = torch.utils.data.RandomSampler(dataset_train)
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.updetr_collate_fn,
num_workers=args.num_workers)
print(len(data_loader_train) * args.epochs)
output_dir = Path(args.output_dir)
if args.resume:
checkpoint = torch.load(args.resume, map_location='cpu')
model_without_ddp.load_state_dict(checkpoint['model'])
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
if lr_scheduler.step_size != args.lr_drop:
lr_scheduler.step_size = args.lr_drop
args.start_epoch = checkpoint['epoch'] + 1
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 20 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 20 == 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 args.output_dir and utils.is_main_process():
with (output_dir / "log.txt").open("a") as f:
f.write(json.dumps(log_stats) + "\n")
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
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print("number of params:", n_parameters)
param_dicts = [
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" not in n and p.requires_grad
]
},
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" in n and p.requires_grad
],
"lr":
args.lr_backbone,
},
]
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)
dataset_train = build_dataset(image_set="train", args=args)
dataset_val = build_dataset(image_set="val", args=args)
if args.distributed:
sampler_train = DistributedSampler(dataset_train)
sampler_val = 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,
)
data_loader_val = DataLoader(
dataset_val,
args.batch_size if args.batch_size < 4 else 4,
sampler=sampler_val,
drop_last=False,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
)
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)
elif args.dataset_file in ["cmdd", "cmdc", "wider"]:
base_ds = None
elif args.dataset_file == "MOT17":
base_ds = dataset_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")
# NOTE: this is Bruno's hack to load stuff in
model_dict = model_without_ddp.state_dict()
pretrained_dict = checkpoint["model"]
# hack for adding query stuff
if ("query_embed.query_embed.weight" in model_dict.keys()
and "query_embed.weight" in pretrained_dict.keys()):
pretrained_dict[
"query_embed.query_embed.weight"] = pretrained_dict[
"query_embed.weight"]
# 1. filter out unnecessary keys
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
# if finetuning skip the linear stuff
if args.finetune:
pretrained_dict = {
k: v
for k, v in pretrained_dict.items()
if k not in ["class_embed.weight", "class_embed.bias"]
}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load new state dict
model_without_ddp.load_state_dict(model_dict)
if (not args.eval and not args.load_model_only
and "optimizer" in checkpoint and "lr_scheduler" in checkpoint
and "epoch" in checkpoint):
optimizer.load_state_dict(checkpoint["optimizer"])
lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
args.start_epoch = checkpoint["epoch"] + 1
if args.eval:
if args.test and args.dataset_file == "wider":
if args.resume:
s = args.resume.split("/")[:-1]
output_dir = "/" + os.path.join(*s)
else:
output_dir = args.output_dir
print("SAVING TEST WIDER TO ", output_dir)
test_wider(
model,
criterion,
postprocessors,
dataset_val,
data_loader_val,
device,
output_dir,
)
return
test_stats, coco_evaluator = evaluate(
model,
criterion,
postprocessors,
data_loader_val,
base_ds,
device,
args.output_dir,
dset_file=args.dataset_file,
)
if args.output_dir and coco_evaluator is not None:
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 100 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 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,
dset_file=args.dataset_file,
)
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)
model, criterion, postprocessors = build_model(args)
model.to(device)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
param_dicts = [
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" not in n and p.requires_grad
]
},
{
"params": [
p for n, p in model_without_ddp.named_parameters()
if "backbone" in n and p.requires_grad
],
"lr":
args.lr_backbone,
},
]
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,
gamma=0.9)
dataset_train = build_dataset(image_set='train', args=args)
dataset_val = build_dataset(image_set='val', args=args)
if args.distributed:
sampler_train = DistributedSampler(dataset_train)
sampler_val = 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)
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)
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)
output_dir = output_dir / f"{args.backbone}_{args.transformer_type}"
if args.output_dir:
output_dir.mkdir(parents=True, exist_ok=True)
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_without_ddp.load_state_dict(checkpoint['model'])
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
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")
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 / f'checkpoint_{epoch}.pth']
# extra checkpoint before LR drop and every 100 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 0:
checkpoint_paths.append(output_dir /
f'checkpoint{epoch}_extra.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 save_on_master(*args, **kwargs):
if is_main_process():
torch.save(*args, **kwargs)