def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
model = torch.nn.parallel.deprecated.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
)
return model
def train(cfg, local_rank, distributed, fp16, dllogger):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
if use_amp:
# Initialize mixed-precision training
if fp16:
use_mixed_precision = True
else:
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = "O1" if use_mixed_precision else "O0"
model, optimizer = amp.initialize(model, optimizer, opt_level=amp_opt_level)
if distributed:
if use_apex_ddp:
model = DDP(model, delay_allreduce=True)
else:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader, iters_per_epoch = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
# set the callback function to evaluate and potentially
# early exit each epoch
if cfg.PER_EPOCH_EVAL:
per_iter_callback_fn = functools.partial(
mlperf_test_early_exit,
iters_per_epoch=iters_per_epoch,
tester=functools.partial(test, cfg=cfg, dllogger=dllogger),
model=model,
distributed=distributed,
min_bbox_map=cfg.MIN_BBOX_MAP,
min_segm_map=cfg.MIN_MASK_MAP,
)
else:
per_iter_callback_fn = None
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
use_amp,
cfg,
dllogger,
per_iter_end_callback_fn=per_iter_callback_fn,
)
return model, iters_per_epoch
def train(cfg, local_rank, distributed):
# original = torch.load('/home/zoey/nas/zoey/github/maskrcnn-benchmark/checkpoints/renderpy150000/model_0025000.pth')
#
# new = {"model": original["model"]}
# torch.save(new, '/home/zoey/nas/zoey/github/maskrcnn-benchmark/checkpoints/finetune/model_0000000.pth')
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
# if cfg.MODEL.DEPTH_ON == True:
# model_depth = build_detection_model(cfg)
# device = torch.device(cfg.MODEL.DEVICE)
# model_depth.to(device)
# optimizer_depth = make_optimizer(cfg, model_depth)
# scheduler_depth = make_lr_scheduler(cfg, optimizer_depth)
# model_depth, optimizer_depth = amp.initialize(model_depth, optimizer_depth, opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpointer = DetectronCheckpointer(cfg,
model,
optimizer,
scheduler,
output_dir,
save_to_disk,
logger=None,
isrgb=True,
isdepth=True)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
# extra_checkpoint_data = checkpointer.load('/home/zoey/nas/zoey/github/maskrcnn-benchmark/checkpoints/renderpy150000/model_0025000.pth')
arguments.update(extra_checkpoint_data)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
do_train(model, data_loader, optimizer, scheduler, checkpointer, device,
checkpoint_period, arguments)
return model
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
logger = logging.getLogger("maskrcnn_benchmark.train")
logger.info("The train model: \n {}".format(model))
device = torch.device(cfg.MODEL.DEVICE)
if cfg.SOLVER.USE_SYNC_BN:
model = apex.parallel.convert_syncbn_model(model)
model.to(device)
optimizer = make_optimizer(cfg, model)
model, optimizer = amp.initialize(model, optimizer, opt_level="O0")
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
# model = torch.nn.parallel.DistributedDataParallel(
# model, device_ids=[local_rank], output_device=local_rank,
# # this should be removed if we update BatchNorm stats
# #broadcast_buffers=False,
# )
model = DDP(model, delay_allreduce=True)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
mode=0,
resolution=None,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
data_loader.collate_fn.special_deal = False
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
writer, arch_writer = setup_writer(output_dir, get_rank())
if arch_writer is not None:
arch_writer.write('Genotype: {}\n'.format(cfg.SEARCH.DECODER.CONFIG))
arch_writer.close()
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
writer,
)
return model
def train_with_validation(cfg, local_rank, distributed, test_weights=None):
arguments = {}
arguments["iteration"] = 0
if test_weights:
cfg.MODEL.WEIGHT = test_weights
cfg.SOLVER.MAX_ITER = 0
ignore_labels = (cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES == 0)
# prepare training data
root_path = os.path.expanduser(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "data")))
data_loader, class_ids = make_data_loader(
root_path,
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
ignore_labels=ignore_labels,
)
# overwrite the number of classes by considering the training set
if cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES <= 0: # if we have binary classification or unknown number of classes
cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES = len(class_ids)
# prepare model
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
# prepare optimizer
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
# prepare validation
run_validation_for_model = partial(run_validation, root_path=root_path, cfg=cfg.clone(), class_ids=class_ids, ignore_labels=ignore_labels, distributed=distributed)
# setup checkpointer
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT, use_latest=False if test_weights else True)
arguments.update(extra_checkpoint_data)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
validation_period = cfg.SOLVER.VALIDATION_PERIOD
# start training
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
validation_period,
checkpoint_period,
arguments,
run_validation_for_model)
return model
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg) # 梦开始的地方
device = torch.device(cfg.MODEL.DEVICE) # ！！！！！
model.to(device)
for name, value in model.backbone.body.network.named_children(
): # 冻结主干网络参数
if int(name) > 60:
for param in value.parameters():
param.requires_grad = False
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16" # 这里可以改成float16来加速
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
test_period = cfg.SOLVER.TEST_PERIOD
if test_period > 0:
data_loader_val = make_data_loader(cfg,
is_train=False,
is_distributed=distributed,
is_for_period=True)
else:
data_loader_val = None
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
do_train(
cfg,
model,
data_loader,
data_loader_val,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
test_period,
arguments,
)
return model
def train(cfg, local_rank, distributed, logger):
if is_main_process():
wandb.init(project='scene-graph',
entity='sgg-speaker-listener',
config=cfg.LISTENER)
debug_print(logger, 'prepare training')
model = build_detection_model(cfg)
listener = build_listener(cfg)
speaker_listener = SpeakerListener(model,
listener,
cfg,
is_joint=cfg.LISTENER.JOINT)
if is_main_process():
wandb.watch(listener)
debug_print(logger, 'end model construction')
# modules that should be always set in eval mode
# their eval() method should be called after model.train() is called
eval_modules = (
model.rpn,
model.backbone,
model.roi_heads.box,
)
fix_eval_modules(eval_modules)
# NOTE, we slow down the LR of the layers start with the names in slow_heads
if cfg.MODEL.ROI_RELATION_HEAD.PREDICTOR == "IMPPredictor":
slow_heads = [
"roi_heads.relation.box_feature_extractor",
"roi_heads.relation.union_feature_extractor.feature_extractor",
]
else:
slow_heads = []
# load pretrain layers to new layers
load_mapping = {
"roi_heads.relation.box_feature_extractor":
"roi_heads.box.feature_extractor",
"roi_heads.relation.union_feature_extractor.feature_extractor":
"roi_heads.box.feature_extractor"
}
if cfg.MODEL.ATTRIBUTE_ON:
load_mapping[
"roi_heads.relation.att_feature_extractor"] = "roi_heads.attribute.feature_extractor"
load_mapping[
"roi_heads.relation.union_feature_extractor.att_feature_extractor"] = "roi_heads.attribute.feature_extractor"
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
listener.to(device)
num_gpus = int(
os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
num_batch = cfg.SOLVER.IMS_PER_BATCH
optimizer = make_optimizer(cfg,
model,
logger,
slow_heads=slow_heads,
slow_ratio=10.0,
rl_factor=float(num_batch))
listener_optimizer = make_listener_optimizer(cfg, listener)
scheduler = make_lr_scheduler(cfg, optimizer, logger)
listener_scheduler = None
debug_print(logger, 'end optimizer and schedule')
if cfg.LISTENER.JOINT:
speaker_listener_optimizer = make_speaker_listener_optimizer(
cfg, speaker_listener.speaker, speaker_listener.listener)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
if cfg.LISTENER.JOINT:
speaker_listener, speaker_listener_optimizer = amp.initialize(
speaker_listener, speaker_listener_optimizer, opt_level='O0')
else:
speaker_listener, listener_optimizer = amp.initialize(
speaker_listener, listener_optimizer, opt_level='O0')
#listener, listener_optimizer = amp.initialize(listener, listener_optimizer, opt_level='O0')
#[model, listener], [optimizer, listener_optimizer] = amp.initialize([model, listener], [optimizer, listener_optimizer], opt_level='O1', loss_scale=1)
#model = amp.initialize(model, opt_level='O1')
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
find_unused_parameters=True,
)
listener = torch.nn.parallel.DistributedDataParallel(
listener,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
find_unused_parameters=True,
)
debug_print(logger, 'end distributed')
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
listener_dir = cfg.LISTENER_DIR
save_to_disk = get_rank() == 0
speaker_checkpointer = DetectronCheckpointer(cfg,
model,
optimizer,
scheduler,
output_dir,
save_to_disk,
custom_scheduler=True)
listener_checkpointer = Checkpointer(listener,
optimizer=listener_optimizer,
save_dir=listener_dir,
save_to_disk=save_to_disk,
custom_scheduler=False)
speaker_listener.add_listener_checkpointer(listener_checkpointer)
speaker_listener.add_speaker_checkpointer(speaker_checkpointer)
speaker_listener.load_listener()
speaker_listener.load_speaker(load_mapping=load_mapping)
debug_print(logger, 'end load checkpointer')
train_data_loader = make_data_loader(cfg,
mode='train',
is_distributed=distributed,
start_iter=arguments["iteration"],
ret_images=True)
val_data_loaders = make_data_loader(cfg,
mode='val',
is_distributed=distributed,
ret_images=True)
debug_print(logger, 'end dataloader')
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
if cfg.SOLVER.PRE_VAL:
logger.info("Validate before training")
#output = run_val(cfg, model, listener, val_data_loaders, distributed, logger)
#print('OUTPUT: ', output)
#(sg_loss, img_loss, sg_acc, img_acc) = output
logger.info("Start training")
meters = MetricLogger(delimiter=" ")
max_iter = len(train_data_loader)
start_iter = arguments["iteration"]
start_training_time = time.time()
end = time.time()
print_first_grad = True
listener_loss_func = torch.nn.MarginRankingLoss(margin=1, reduction='none')
mistake_saver = None
if is_main_process():
ds_catalog = DatasetCatalog()
dict_file_path = os.path.join(
ds_catalog.DATA_DIR,
ds_catalog.DATASETS['VG_stanford_filtered_with_attribute']
['dict_file'])
ind_to_classes, ind_to_predicates = load_vg_info(dict_file_path)
ind_to_classes = {k: v for k, v in enumerate(ind_to_classes)}
ind_to_predicates = {k: v for k, v in enumerate(ind_to_predicates)}
print('ind to classes:', ind_to_classes, '/n ind to predicates:',
ind_to_predicates)
mistake_saver = MistakeSaver(
'/Scene-Graph-Benchmark.pytorch/filenames_masked', ind_to_classes,
ind_to_predicates)
#is_printed = False
while True:
try:
listener_iteration = 0
for iteration, (images, targets,
image_ids) in enumerate(train_data_loader,
start_iter):
if cfg.LISTENER.JOINT:
speaker_listener_optimizer.zero_grad()
else:
listener_optimizer.zero_grad()
#print(f'ITERATION NUMBER: {iteration}')
if any(len(target) < 1 for target in targets):
logger.error(
f"Iteration={iteration + 1} || Image Ids used for training {_} || targets Length={[len(target) for target in targets]}"
)
if len(images) <= 1:
continue
data_time = time.time() - end
iteration = iteration + 1
listener_iteration += 1
arguments["iteration"] = iteration
model.train()
fix_eval_modules(eval_modules)
images_list = deepcopy(images)
images_list = to_image_list(
images_list, cfg.DATALOADER.SIZE_DIVISIBILITY).to(device)
for i in range(len(images)):
images[i] = images[i].unsqueeze(0)
images[i] = F.interpolate(images[i],
size=(224, 224),
mode='bilinear',
align_corners=False)
images[i] = images[i].squeeze()
images = torch.stack(images).to(device)
#images.requires_grad_()
targets = [target.to(device) for target in targets]
speaker_loss_dict = {}
if not cfg.LISTENER.JOINT:
score_matrix = speaker_listener(images_list, targets,
images)
else:
score_matrix, _, speaker_loss_dict = speaker_listener(
images_list, targets, images)
speaker_summed_losses = sum(
loss for loss in speaker_loss_dict.values())
# reduce losses over all GPUs for logging purposes
if not not cfg.LISTENER.JOINT:
speaker_loss_dict_reduced = reduce_loss_dict(
speaker_loss_dict)
speaker_losses_reduced = sum(
loss for loss in speaker_loss_dict_reduced.values())
speaker_losses_reduced /= num_gpus
if is_main_process():
wandb.log(
{"Train Speaker Loss": speaker_losses_reduced},
listener_iteration)
listener_loss = 0
gap_reward = 0
avg_acc = 0
num_correct = 0
score_matrix = score_matrix.to(device)
# fill loss matrix
loss_matrix = torch.zeros((2, images.size(0), images.size(0)),
device=device)
# sg centered scores
for true_index in range(loss_matrix.size(1)):
row_score = score_matrix[true_index]
(true_scores, predicted_scores,
binary) = format_scores(row_score, true_index, device)
loss_vec = listener_loss_func(true_scores,
predicted_scores, binary)
loss_matrix[0][true_index] = loss_vec
# image centered scores
transposted_score_matrix = score_matrix.t()
for true_index in range(loss_matrix.size(1)):
row_score = transposted_score_matrix[true_index]
(true_scores, predicted_scores,
binary) = format_scores(row_score, true_index, device)
loss_vec = listener_loss_func(true_scores,
predicted_scores, binary)
loss_matrix[1][true_index] = loss_vec
print('iteration:', listener_iteration)
sg_acc = 0
img_acc = 0
# calculate accuracy
for i in range(loss_matrix.size(1)):
temp_sg_acc = 0
temp_img_acc = 0
for j in range(loss_matrix.size(2)):
if loss_matrix[0][i][i] > loss_matrix[0][i][j]:
temp_sg_acc += 1
else:
if cfg.LISTENER.HTML:
if is_main_process(
) and listener_iteration >= 600 and listener_iteration % 25 == 0 and i != j:
detached_sg_i = (sgs[i][0].detach(),
sgs[i][1],
sgs[i][2].detach())
detached_sg_j = (sgs[j][0].detach(),
sgs[j][1],
sgs[j][2].detach())
mistake_saver.add_mistake(
(image_ids[i], image_ids[j]),
(detached_sg_i, detached_sg_j),
listener_iteration, 'SG')
if loss_matrix[1][i][i] > loss_matrix[1][j][i]:
temp_img_acc += 1
else:
if cfg.LISTENER.HTML:
if is_main_process(
) and listener_iteration >= 600 and listener_iteration % 25 == 0 and i != j:
detached_sg_i = (sgs[i][0].detach(),
sgs[i][1],
sgs[i][2].detach())
detached_sg_j = (sgs[j][0].detach(),
sgs[j][1],
sgs[j][2].detach())
mistake_saver.add_mistake(
(image_ids[i], image_ids[j]),
(detached_sg_i, detached_sg_j),
listener_iteration, 'IMG')
temp_sg_acc = temp_sg_acc * 100 / (loss_matrix.size(1) - 1)
temp_img_acc = temp_img_acc * 100 / (loss_matrix.size(1) -
1)
sg_acc += temp_sg_acc
img_acc += temp_img_acc
if cfg.LISTENER.HTML:
if is_main_process(
) and listener_iteration % 100 == 0 and listener_iteration >= 600:
mistake_saver.toHtml('/www')
sg_acc /= loss_matrix.size(1)
img_acc /= loss_matrix.size(1)
avg_sg_acc = torch.tensor([sg_acc]).to(device)
avg_img_acc = torch.tensor([img_acc]).to(device)
# reduce acc over all gpus
avg_acc = {'sg_acc': avg_sg_acc, 'img_acc': avg_img_acc}
avg_acc_reduced = reduce_loss_dict(avg_acc)
sg_acc = sum(acc for acc in avg_acc_reduced['sg_acc'])
img_acc = sum(acc for acc in avg_acc_reduced['img_acc'])
# log acc to wadb
if is_main_process():
wandb.log({
"Train SG Accuracy": sg_acc.item(),
"Train IMG Accuracy": img_acc.item()
})
sg_loss = 0
img_loss = 0
for i in range(loss_matrix.size(0)):
for j in range(loss_matrix.size(1)):
loss_matrix[i][j][j] = 0.
for i in range(loss_matrix.size(1)):
sg_loss += torch.max(loss_matrix[0][i])
img_loss += torch.max(loss_matrix[1][:][i])
sg_loss = sg_loss / loss_matrix.size(1)
img_loss = img_loss / loss_matrix.size(1)
sg_loss = sg_loss.to(device)
img_loss = img_loss.to(device)
loss_dict = {'sg_loss': sg_loss, 'img_loss': img_loss}
losses = sum(loss for loss in loss_dict.values())
# reduce losses over all GPUs for logging purposes
loss_dict_reduced = reduce_loss_dict(loss_dict)
sg_loss_reduced = loss_dict_reduced['sg_loss']
img_loss_reduced = loss_dict_reduced['img_loss']
if is_main_process():
wandb.log({"Train SG Loss": sg_loss_reduced})
wandb.log({"Train IMG Loss": img_loss_reduced})
losses_reduced = sum(loss
for loss in loss_dict_reduced.values())
meters.update(loss=losses_reduced, **loss_dict_reduced)
losses = losses + speaker_summed_losses * cfg.LISTENER.LOSS_COEF
# Note: If mixed precision is not used, this ends up doing nothing
# Otherwise apply loss scaling for mixed-precision recipe
#losses.backward()
if not cfg.LISTENER.JOINT:
with amp.scale_loss(losses,
listener_optimizer) as scaled_losses:
scaled_losses.backward()
else:
with amp.scale_loss(
losses,
speaker_listener_optimizer) as scaled_losses:
scaled_losses.backward()
verbose = (iteration % cfg.SOLVER.PRINT_GRAD_FREQ
) == 0 or print_first_grad # print grad or not
print_first_grad = False
#clip_grad_value([(n, p) for n, p in listener.named_parameters() if p.requires_grad], cfg.LISTENER.CLIP_VALUE, logger=logger, verbose=True, clip=True)
if not cfg.LISTENER.JOINT:
listener_optimizer.step()
else:
speaker_listener_optimizer.step()
batch_time = time.time() - end
end = time.time()
meters.update(time=batch_time, data=data_time)
eta_seconds = meters.time.global_avg * (max_iter - iteration)
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
if cfg.LISTENER.JOINT:
if iteration % 200 == 0 or iteration == max_iter:
logger.info(
meters.delimiter.join([
"eta: {eta}",
"iter: {iter}",
"{meters}",
"lr: {lr:.6f}",
"max mem: {memory:.0f}",
]).format(
eta=eta_string,
iter=iteration,
meters=str(meters),
lr=speaker_listener_optimizer.param_groups[-1]
["lr"],
memory=torch.cuda.max_memory_allocated() /
1024.0 / 1024.0,
))
else:
if iteration % 200 == 0 or iteration == max_iter:
logger.info(
meters.delimiter.join([
"eta: {eta}",
"iter: {iter}",
"{meters}",
"lr: {lr:.6f}",
"max mem: {memory:.0f}",
]).format(
eta=eta_string,
iter=iteration,
meters=str(meters),
lr=listener_optimizer.param_groups[-1]["lr"],
memory=torch.cuda.max_memory_allocated() /
1024.0 / 1024.0,
))
if iteration % checkpoint_period == 0:
"""
print('Model before save')
print('****************************')
print(listener.gnn.conv1.node_model.node_mlp_1[0].weight)
print('****************************')
"""
if not cfg.LISTENER.JOINT:
listener_checkpointer.save(
"model_{:07d}".format(listener_iteration),
amp=amp.state_dict())
else:
speaker_checkpointer.save(
"model_speaker{:07d}".format(iteration))
listener_checkpointer.save(
"model_listenr{:07d}".format(listener_iteration),
amp=amp.state_dict())
if iteration == max_iter:
if not cfg.LISTENER.JOINT:
listener_checkpointer.save(
"model_{:07d}".format(listener_iteration),
amp=amp.state_dict())
else:
speaker_checkpointer.save(
"model_{:07d}".format(iteration))
listener_checkpointer.save(
"model_{:07d}".format(listener_iteration),
amp=amp.state_dict())
val_result = None # used for scheduler updating
if cfg.SOLVER.TO_VAL and iteration % cfg.SOLVER.VAL_PERIOD == 0:
logger.info("Start validating")
val_result = run_val(cfg, model, listener,
val_data_loaders, distributed, logger)
(sg_loss, img_loss, sg_acc, img_acc,
speaker_val) = val_result
if is_main_process():
wandb.log({
"Validation SG Accuracy": sg_acc,
"Validation IMG Accuracy": img_acc,
"Validation SG Loss": sg_loss,
"Validation IMG Loss": img_loss,
"Validation Speaker": speaker_val,
})
#logger.info("Validation Result: %.4f" % val_result)
except Exception as err:
raise (err)
print('Dataset finished, creating new')
train_data_loader = make_data_loader(
cfg,
mode='train',
is_distributed=distributed,
start_iter=arguments["iteration"],
ret_images=True)
total_training_time = time.time() - start_training_time
total_time_str = str(datetime.timedelta(seconds=total_training_time))
logger.info("Total training time: {} ({:.4f} s / it)".format(
total_time_str, total_training_time / (max_iter)))
return listener
def train(cfg, local_rank, distributed):
# Model logging
print_mlperf(key=mlperf_log.INPUT_BATCH_SIZE, value=cfg.SOLVER.IMS_PER_BATCH)
print_mlperf(key=mlperf_log.BATCH_SIZE_TEST, value=cfg.TEST.IMS_PER_BATCH)
print_mlperf(key=mlperf_log.INPUT_MEAN_SUBTRACTION, value = cfg.INPUT.PIXEL_MEAN)
print_mlperf(key=mlperf_log.INPUT_NORMALIZATION_STD, value=cfg.INPUT.PIXEL_STD)
print_mlperf(key=mlperf_log.INPUT_RESIZE)
print_mlperf(key=mlperf_log.INPUT_RESIZE_ASPECT_PRESERVING)
print_mlperf(key=mlperf_log.MIN_IMAGE_SIZE, value=cfg.INPUT.MIN_SIZE_TRAIN)
print_mlperf(key=mlperf_log.MAX_IMAGE_SIZE, value=cfg.INPUT.MAX_SIZE_TRAIN)
print_mlperf(key=mlperf_log.INPUT_RANDOM_FLIP)
print_mlperf(key=mlperf_log.RANDOM_FLIP_PROBABILITY, value=0.5)
print_mlperf(key=mlperf_log.FG_IOU_THRESHOLD, value=cfg.MODEL.RPN.FG_IOU_THRESHOLD)
print_mlperf(key=mlperf_log.BG_IOU_THRESHOLD, value=cfg.MODEL.RPN.BG_IOU_THRESHOLD)
print_mlperf(key=mlperf_log.RPN_PRE_NMS_TOP_N_TRAIN, value=cfg.MODEL.RPN.PRE_NMS_TOP_N_TRAIN)
print_mlperf(key=mlperf_log.RPN_PRE_NMS_TOP_N_TEST, value=cfg.MODEL.RPN.PRE_NMS_TOP_N_TEST)
print_mlperf(key=mlperf_log.RPN_POST_NMS_TOP_N_TRAIN, value=cfg.MODEL.RPN.FPN_POST_NMS_TOP_N_TRAIN)
print_mlperf(key=mlperf_log.RPN_POST_NMS_TOP_N_TEST, value=cfg.MODEL.RPN.FPN_POST_NMS_TOP_N_TEST)
print_mlperf(key=mlperf_log.ASPECT_RATIOS, value=cfg.MODEL.RPN.ASPECT_RATIOS)
print_mlperf(key=mlperf_log.BACKBONE, value=cfg.MODEL.BACKBONE.CONV_BODY)
print_mlperf(key=mlperf_log.NMS_THRESHOLD, value=cfg.MODEL.RPN.NMS_THRESH)
# /root/ssy/ssynew/maskrcnn-benchmark/maskrcnn_benchmark/modeling/detector/detectors.py
# building bare mode without doing anthing
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
# Optimizer logging
print_mlperf(key=mlperf_log.OPT_NAME, value=mlperf_log.SGD_WITH_MOMENTUM)
print_mlperf(key=mlperf_log.OPT_LR, value=cfg.SOLVER.BASE_LR)
print_mlperf(key=mlperf_log.OPT_MOMENTUM, value=cfg.SOLVER.MOMENTUM)
print_mlperf(key=mlperf_log.OPT_WEIGHT_DECAY, value=cfg.SOLVER.WEIGHT_DECAY)
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
print("output_dir "+str(output_dir))
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
# no such SAVE_CHECKPOINTS
#arguments["save_checkpoints"] = cfg.SAVE_CHECKPOINTS
arguments["save_checkpoints"] = False
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader, iters_per_epoch = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"]
)
print("SSY iters_per_epoch "+str(iters_per_epoch))
#print("SSY iters_per_epoch change to 100 ")
#iters_per_epoch = 100
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
# set the callback function to evaluate and potentially
# early exit each epoch
# SSY
# I already add PER_EPOCH_EVAL and MIN_BBOX_MAP MIN_SEGM_MAP to ./configs/e2e_mask_rcnn_R_50_FPN_1x.yaml
# but it still can not find it
# so I manually set them here
#if cfg.PER_EPOCH_EVAL:
# per_iter_callback_fn = functools.partial(
# mlperf_test_early_exit,
# iters_per_epoch=iters_per_epoch,
# tester=functools.partial(test, cfg=cfg),
# model=model,
# distributed=distributed,
# min_bbox_map=cfg.MLPERF.MIN_BBOX_MAP,
# min_segm_map=cfg.MLPERF.MIN_SEGM_MAP)
#else:
# per_iter_callback_fn = None
per_iter_callback_fn = functools.partial(
mlperf_test_early_exit,
iters_per_epoch=iters_per_epoch,
# /root/ssy/ssynew/maskrcnn-benchmark/maskrcnn_benchmark/engine/tester.py
tester=functools.partial(test, cfg=cfg),
model=model,
distributed=distributed,
min_bbox_map=0.377,
min_segm_map=0.339)
start_train_time = time.time()
# /root/ssy/ssynew/maskrcnn-benchmark/maskrcnn_benchmark/engine/trainer.py
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
per_iter_start_callback_fn=functools.partial(mlperf_log_epoch_start, iters_per_epoch=iters_per_epoch),
per_iter_end_callback_fn=per_iter_callback_fn,
)
end_train_time = time.time()
total_training_time = end_train_time - start_train_time
print(
"&&&& MLPERF METRIC THROUGHPUT per GPU={:.4f} iterations / s".format((arguments["iteration"] * 1.0) / total_training_time)
)
return model
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
# pdb.set_trace()
if cfg.MODEL.USE_SYNCBN:
assert is_pytorch_1_1_0_or_later(), \
"SyncBatchNorm is only available in pytorch >= 1.1.0"
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# pdb.set_trace()
# (Pdb) optimizer.param_groups[0]["lr"]
# 0.0016666666666666666
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
# pdb.set_trace()
# (Pdb) optimizer.param_groups[0]["lr"]
# 0.0016666666666666666
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
# pdb.set_trace()
# (Pdb) cfg.MODEL.WEIGHT
# 'coco_P2_8.pth'
# (Pdb) optimizer.param_groups[0]["lr"]
# 0.0016666666666666666
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
# pdb.set_trace()
# (Pdb) optimizer.param_groups[0]["lr"]
# 0.00010000000000000002
# pdb.set_trace()
# (Pdb) extra_checkpoint_data
# {'iteration': 80000}
# (Pdb) arguments
# {'iteration': 80000}
# coco_pretrained_P2, start=8000 => start=0
arguments["iteration"] = 0
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
# pdb.set_trace()
# (Pdb) cfg.SOLVER.CHECKPOINT_PERIOD
# 2500
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
# pdb.set_trace()
# optimizer.param_groups[0]["lr"]
# 0.00010000000000000002
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
)
return model
def train(cfg, args):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
if use_amp:
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model, optimizer, opt_level=amp_opt_level)
if args.distributed:
# if use_apex_ddp:
# model = DDP(model, delay_allreduce=True)
# else:
# SMDataParallel: Wrap the PyTorch model with SMDataParallel’s DDP
model = DDP(model, device_ids=[dist.get_local_rank()], broadcast_buffers=False)
#model = DDP(model)
print("model parameter size: ", sum(p.numel() for p in model.parameters() if p.requires_grad))
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
# SMDataParallel: Save model on master node.
save_to_disk = dist.get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader, iters_per_epoch = make_data_loader(
cfg,
is_train=True,
is_distributed=args.distributed,
start_iter=arguments["iteration"],
data_dir = args.data_dir
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
# set the callback function to evaluate and potentially
# early exit each epoch
if cfg.PER_EPOCH_EVAL:
per_iter_callback_fn = functools.partial(
mlperf_test_early_exit,
iters_per_epoch=iters_per_epoch,
tester=functools.partial(test, cfg=cfg),
model=model,
distributed=args.distributed,
min_bbox_map=cfg.MIN_BBOX_MAP,
min_segm_map=cfg.MIN_MASK_MAP)
else:
per_iter_callback_fn = None
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
use_amp,
cfg,
per_iter_end_callback_fn=per_iter_callback_fn,
)
return model
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
# pickle.load = partial(pickle.load, encoding="latin1")
# pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1")
# pretrained_dict= torch.load(cfg.MODEL.WEIGHT, map_location=lambda storage, loc: storage, pickle_module=pickle)
# # pretrained_dict=torch.load(cfg.MODEL.WEIGHT)
# model_dict=model.state_dict()
# pretrained_dict={k: v for k, v in pretrained_dict.items() if k in model_dict}
# model_dict.update(pretrained_dict)
# # torch.save(model_dict,'./pretrained.pkl')
# model.load_state_dict(model_dict)
# extra_checkpoint_data = checkpointer.load('./pretrained.pkl',use_latest=False)
# arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
)
return model
def train(cfg, local_rank, distributed, random_number_generator=None):
if (torch._C, '_jit_set_profiling_executor'):
torch._C._jit_set_profiling_executor(False)
if (torch._C, '_jit_set_profiling_mode'):
torch._C._jit_set_profiling_mode(False)
# Model logging
log_event(key=constants.GLOBAL_BATCH_SIZE, value=cfg.SOLVER.IMS_PER_BATCH)
log_event(key=constants.NUM_IMAGE_CANDIDATES,
value=cfg.MODEL.RPN.FPN_POST_NMS_TOP_N_TRAIN)
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
print("DEVICE IS : {}".format(device))
# Initialize mixed-precision training
is_fp16 = (cfg.DTYPE == "float16")
if is_fp16:
# convert model to FP16
model.half()
# Optimizer logging
log_event(key=constants.OPT_NAME, value="sgd_with_momentum")
log_event(key=constants.OPT_BASE_LR, value=cfg.SOLVER.BASE_LR)
log_event(key=constants.OPT_LR_WARMUP_STEPS, value=cfg.SOLVER.WARMUP_ITERS)
log_event(key=constants.OPT_LR_WARMUP_FACTOR,
value=cfg.SOLVER.WARMUP_FACTOR)
log_event(key=constants.OPT_LR_DECAY_FACTOR, value=cfg.SOLVER.GAMMA)
log_event(key=constants.OPT_LR_DECAY_STEPS, value=cfg.SOLVER.STEPS)
log_event(key=constants.MIN_IMAGE_SIZE, value=cfg.INPUT.MIN_SIZE_TRAIN[0])
log_event(key=constants.MAX_IMAGE_SIZE, value=cfg.INPUT.MAX_SIZE_TRAIN)
scheduler = make_lr_scheduler(cfg, optimizer)
# disable the garbage collection
gc.disable()
if distributed:
model = DDP(model,
device_ids=[herring.get_local_rank()],
broadcast_buffers=False)
arguments = {}
arguments["iteration"] = 0
arguments["nhwc"] = cfg.NHWC
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
arguments["save_checkpoints"] = cfg.SAVE_CHECKPOINTS
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT, cfg.NHWC)
arguments.update(extra_checkpoint_data)
if is_fp16:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
log_end(key=constants.INIT_STOP)
barrier()
log_start(key=constants.RUN_START)
barrier()
data_loader, iters_per_epoch = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
random_number_generator=random_number_generator,
)
log_event(key=constants.TRAIN_SAMPLES, value=len(data_loader))
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
# set the callback function to evaluate and potentially
# early exit each epoch
if cfg.PER_EPOCH_EVAL:
per_iter_callback_fn = functools.partial(
mlperf_test_early_exit,
iters_per_epoch=iters_per_epoch,
tester=functools.partial(test, cfg=cfg),
model=model,
distributed=distributed,
min_bbox_map=cfg.MLPERF.MIN_BBOX_MAP,
min_segm_map=cfg.MLPERF.MIN_SEGM_MAP)
else:
per_iter_callback_fn = None
start_train_time = time.time()
success = do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
cfg.DISABLE_REDUCED_LOGGING,
per_iter_start_callback_fn=functools.partial(
mlperf_log_epoch_start, iters_per_epoch=iters_per_epoch),
per_iter_end_callback_fn=per_iter_callback_fn,
)
end_train_time = time.time()
total_training_time = end_train_time - start_train_time
print("&&&& MLPERF METRIC THROUGHPUT={:.4f} iterations / s".format(
(arguments["iteration"] * cfg.SOLVER.IMS_PER_BATCH) /
total_training_time))
return model, success
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
# NOTE Deyu
# 加載dota完整預訓練權重，使用自定義權重管理類
# checkpointer = TransferLearningCheckpointer(
# cfg, model, optimizer, scheduler, output_dir, save_to_disk
# )
# extra_checkpoint_data = checkpointer.load_checkpoint_pop(cfg.MODEL.WEIGHT)
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
# NOTE Mingtao
lr_mingtao = scheduler.base_lrs
# NOTE Mingtao: force to use new steps
scheduler.milestones = cfg.SOLVER.STEPS
scheduler.base_lrs = lr_mingtao
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
checkpoint_start_step = cfg.SOLVER.CHECKPOINT_START_STEP
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
checkpoint_start_step,
arguments,
)
return model
def train(cfg, local_rank, distributed):
model = create_model(cfg)
model_ema = create_model(cfg, ema=True)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
model_ema.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
model_ema = DistributedDataParallel(model_ema)
arguments = {}
cfg_arg = {}
arguments["iteration"] = 0
arguments["semi_weight"] = cfg.SEMI.SEMI_WEIGHT
cfg_arg["temporal_save_path"] = cfg.SEMI.TEMPORAL_SAVE_PATH
arguments['loss_semi'] = make_semi_box_loss_evaluator(cfg)
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader_semi(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
update_ema_variables(model, model_ema)
# arguments["iteration"] = 0
# optimizer = make_optimizer(cfg, model)
# scheduler = make_lr_scheduler(cfg, optimizer)
arguments["ema_decay"] = cfg.SEMI.EMA_DECAY
arguments["ANCHOR_STRIDES"] = cfg.MODEL.RETINANET.ANCHOR_STRIDES
arguments["HYPER_PARAMETERS"] = cfg.SEMI.HYPER_PARAMETERS
arguments['postprocess'] = make_retinanet_semi_postprocessor(
cfg, BoxCoder(weights=(10., 10., 5., 5.)), True)
for g in optimizer.param_groups:
g['lr'] = 0.0005
do_train(
model,
model_ema,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
local_rank,
checkpoint_period,
cfg_arg,
arguments,
)
return model
def train(cfg, local_rank, distributed):
# 创建GeneralizedRCNN()对象
# detectors.py --> generalized_rcnn.py
model = build_detection_model(cfg)
# 'cpu' or 'cuda'
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
# 封装了 torch.optiom.SGD() 函数, 根据tensor的requires_grad属性构成需要更新的参数列表
optimizer = make_optimizer(cfg, model)
# 根据配置信息设置 optimizer 的学习率更新策略
scheduler = make_lr_scheduler(cfg, optimizer)
# 分布式训练情况下, 并行处理数据
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
# 获取输出的文件夹路径, 默认为 '.', 配置文件中设置为'./log'
output_dir = cfg.OUTPUT_DIR
# 如果分布式训练不可用, 则将这个变量设置为True
save_to_disk = get_rank() == 0
checkpointer = \
DetectronCheckpointer(cfg, model, optimizer, scheduler, output_dir, save_to_disk)
# cfg.MODEL.WEIGHT="catalog://ImageNetPretrained/MSRA/R-50"
# 这个实际上是个空字典, 预训练模型中只有'model'一个key
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD # 2500
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
)
return model
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# if use_amp:
# # Initialize mixed-precision training
# use_mixed_precision = cfg.DTYPE == "float16"
# amp_handle = amp.init(enabled=use_mixed_precision, verbose=cfg.AMP_VERBOSE)
# # wrap the optimizer for mixed precision
# if cfg.SOLVER.ACCUMULATE_GRAD:
# # also specify number of steps to accumulate over
# optimizer = amp_handle.wrap_optimizer(optimizer, num_loss=cfg.SOLVER.ACCUMULATE_STEPS)
# else:
# optimizer = amp_handle.wrap_optimizer(optimizer)
model, optimizer = amp.initialize(model, optimizer,opt_level='O1')
if distributed:
if use_apex_ddp:
model = DDP(model, delay_allreduce=True)
else:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader, iters_per_epoch = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
# set the callback function to evaluate and potentially
# early exit each epoch
if 1==1:
per_iter_callback_fn = functools.partial(
mlperf_test_early_exit,
iters_per_epoch=iters_per_epoch,
tester=functools.partial(test, cfg=cfg),
model=model,
distributed=distributed,
min_bbox_map=cfg.MIN_BBOX_MAP,
min_segm_map=cfg.MIN_MASK_MAP)
else:
per_iter_callback_fn = None
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
use_amp,
cfg,
per_iter_end_callback_fn=per_iter_callback_fn,
)
return model
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
params = get_model_parameters_number(model)
print('{:<30} {:<8}'.format('Number of parameters: ', params))
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
if cfg.MODEL.DOMAIN_ADAPTATION_ON:
source_data_loader = make_data_loader(
cfg,
is_train=True,
is_source=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
target_data_loader = make_data_loader(
cfg,
is_train=True,
is_source=False,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
do_da_train(
model,
source_data_loader,
target_data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
cfg,
)
else:
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
)
return model
def train(cfg, cfg_origial, local_rank, distributed):
## The one with modified number of classes
model = build_detection_model(cfg)
# cfg_origial = cfg.clone()
# cfg_origial.MODEL.ROI_BOX_HEAD.NUM_CLASSES = 81
# original_model = build_detection_model(cfg_origial) ## Original model with 81 classes
# ## Let's load weights for old class!
# save_dir = cfg.OUTPUT_DIR
# checkpointer = DetectronCheckpointer(cfg_origial, original_model, save_dir=save_dir)
# checkpointer.load(cfg_origial.MODEL.WEIGHT)
# # pretrained_model_pth = "/network/home/bhattdha/.torch/models/_detectron_35861795_12_2017_baselines_e2e_mask_rcnn_R-101-FPN_1x.yaml.02_31_37.KqyEK4tT_output_train_coco_2014_train%3Acoco_2014_valminusminival_generalized_rcnn_model_final.pkl"
# # These keys are to be removed which forms final layers of the network
# removal_keys = ['roi_heads.box.predictor.cls_score.weight', 'roi_heads.box.predictor.cls_score.bias', 'roi_heads.box.predictor.bbox_pred.weight', 'roi_heads.box.predictor.bbox_pred.bias', 'roi_heads.mask.predictor.mask_fcn_logits.weight', 'roi_heads.mask.predictor.mask_fcn_logits.bias']
# model = _transfer_pretrained_weights(new_model, original_model, removal_keys)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# # Initialize mixed-precision training
# use_mixed_precision = cfg.DTYPE == "float16"
# amp_opt_level = 'O1' if use_mixed_precision else 'O0'
# model, optimizer = amp.initialize(model, optimizer, opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
# cfg.MODEL.WEIGHT = '/network/home/bhattdha/exp.pth' ## Model stored through surgery
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
)
return model
def train(self, is_train, result_dir=None):
model = build_detection_model(self.cfg)
device = torch.device(self.cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(self.cfg, model)
scheduler = make_lr_scheduler(self.cfg, optimizer)
# Initialize mixed-precision training
use_mixed_precision = self.cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if self.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[self.local_rank],
output_device=self.local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = self.cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(self.cfg, model, optimizer,
scheduler, output_dir,
save_to_disk)
# Load rpn
if self.cfg.MODEL.WEIGHT.startswith(
'/') or 'catalog' in self.cfg.MODEL.WEIGHT:
model_path = self.cfg.MODEL.WEIGHT
else:
model_path = os.path.abspath(
os.path.join(os.path.dirname(__file__), os.path.pardir,
os.path.pardir, os.path.pardir, os.path.pardir,
'Data', 'pretrained_feature_extractors',
self.cfg.MODEL.WEIGHT))
checkpointer = DetectronCheckpointer(cfg, model, save_dir=result_dir)
_ = checkpointer.load(model_path)
if self.distributed:
model = model.module
iou_types = ("bbox", )
torch.cuda.empty_cache() # TODO check if it helps
output_folders = [None]
if is_train:
dataset_names = ['train']
else:
dataset_names = ['test']
if self.cfg.OUTPUT_DIR:
for idx, dataset_name in enumerate(dataset_names):
output_folder = os.path.join(self.cfg.OUTPUT_DIR, dataset_name)
mkdir(output_folder)
output_folders[idx] = output_folder
data_loaders = make_data_loader(self.cfg,
is_train=is_train,
is_distributed=self.distributed,
is_final_test=True,
is_target_task=self.is_target_task,
icwt_21_objs=self.icwt_21_objs)
for output_folder, dataset_name, data_loader in zip(
output_folders, dataset_names, data_loaders):
feat_extraction_time = inference(
self.cfg,
model,
data_loader,
dataset_name=dataset_name,
iou_types=iou_types,
box_only=False
if self.cfg.MODEL.RETINANET_ON else self.cfg.MODEL.RPN_ONLY,
device=cfg.MODEL.DEVICE,
is_target_task=self.is_target_task,
icwt_21_objs=self.icwt_21_objs,
is_train=is_train,
result_dir=result_dir,
)
if result_dir and is_train:
with open(os.path.join(result_dir, "result.txt"), "a") as fid:
fid.write(
"RPN's feature extraction time: {}min:{}s \n".format(
int(feat_extraction_time / 60),
round(feat_extraction_time % 60)))
synchronize()
logger = logging.getLogger("maskrcnn_benchmark")
logger.handlers = []
if self.cfg.SAVE_FEATURES_RPN:
# Save features still not saved
for clss in model.rpn.anchors_ids:
# Save negatives batches
for batch in range(len(model.rpn.negatives[clss])):
if model.rpn.negatives[clss][batch].size()[0] > 0:
path_to_save = os.path.join(
result_dir, 'features_RPN',
'negatives_cl_{}_batch_{}'.format(clss, batch))
torch.save(model.rpn.negatives[clss][batch],
path_to_save)
# If a class does not have positive examples, save an empty tensor
if model.rpn.positives[clss][0].size()[0] == 0 and len(
model.rpn.positives[clss]) == 1:
path_to_save = os.path.join(
result_dir, 'features_RPN',
'positives_cl_{}_batch_{}'.format(clss, 0))
torch.save(
torch.empty(
(0, model.rpn.feat_size),
device=model.rpn.negatives[clss][0].device),
path_to_save)
else:
for batch in range(len(model.rpn.positives[clss])):
if model.rpn.positives[clss][batch].size()[0] > 0:
path_to_save = os.path.join(
result_dir, 'features_RPN',
'positives_cl_{}_batch_{}'.format(clss, batch))
torch.save(model.rpn.positives[clss][batch],
path_to_save)
for i in range(len(model.rpn.X)):
if model.rpn.X[i].size()[0] > 0:
path_to_save = os.path.join(result_dir, 'features_RPN',
'reg_x_batch_{}'.format(i))
torch.save(model.rpn.X[i], path_to_save)
path_to_save = os.path.join(result_dir, 'features_RPN',
'reg_c_batch_{}'.format(i))
torch.save(model.rpn.C[i], path_to_save)
path_to_save = os.path.join(result_dir, 'features_RPN',
'reg_y_batch_{}'.format(i))
torch.save(model.rpn.Y[i], path_to_save)
return
else:
COXY = {
'C': torch.cat(model.rpn.C),
'O': model.rpn.O,
'X': torch.cat(model.rpn.X),
'Y': torch.cat(model.rpn.Y)
}
for i in range(self.cfg.MINIBOOTSTRAP.RPN.NUM_CLASSES):
model.rpn.positives[i] = torch.cat(model.rpn.positives[i])
return copy.deepcopy(model.rpn.negatives), copy.deepcopy(
model.rpn.positives), copy.deepcopy(COXY)
def train(cfg, local_rank, distributed, use_tensorboard=False, logger=None):
arguments = {"iteration": 0}
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
if cfg.SOLVER.UNFREEZE_CONV_BODY:
for p in model.backbone.parameters():
p.requires_grad = True
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg,
model,
optimizer,
scheduler,
output_dir,
save_to_disk,
logger=logger)
print(cfg.TRAIN.IGNORE_LIST)
extra_checkpoint_data = checkpointer.load(
cfg.MODEL.WEIGHT, ignore_list=cfg.TRAIN.IGNORE_LIST)
arguments.update(extra_checkpoint_data)
if cfg.SOLVER.KEEP_LR:
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
tensorboard_logdir = cfg.OUTPUT_DIR
tensorboard_exp_name = cfg.TENSORBOARD_EXP_NAME
snapshot = cfg.SOLVER.SNAPSHOT_ITERS
do_train(model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
snapshot,
tensorboard_logdir,
tensorboard_exp_name,
use_tensorboard=use_tensorboard)
return model
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
summary_writer = SummaryWriter(log_dir=output_dir)
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
if cfg.MODEL.WEIGHT.upper() == 'CONTINUE':
model_weight = last_checkpoint(output_dir)
else:
model_weight = cfg.MODEL.WEIGHT
extra_checkpoint_data = checkpointer.load(model_weight)
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
data_loader_val = make_data_loader(
cfg,
is_train=False,
is_distributed=distributed)[0]
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
do_train(
model=model,
data_loader=data_loader,
data_loader_val=data_loader_val,
optimizer=optimizer,
scheduler=scheduler,
checkpointer=checkpointer,
device=device,
checkpoint_period=checkpoint_period,
arguments=arguments,
summary_writer=summary_writer
)
return model
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
# import ipdb;ipdb.set_trace()
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
if not os.path.exists(output_dir):
os.makedirs(output_dir)
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
test_period = cfg.SOLVER.TEST_PERIOD
if test_period > 0:
data_loader_val = make_data_loader(cfg,
is_train=False,
is_distributed=distributed,
is_for_period=True)
else:
data_loader_val = None
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
do_train(
cfg,
model,
data_loader,
data_loader_val,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
test_period,
arguments,
)
return model
def train(cfg, local_rank, distributed, logger):
debug_print(logger, 'prepare training')
model = build_detection_model(cfg)
debug_print(logger, 'end model construction')
# modules that should be always set in eval mode
# their eval() method should be called after model.train() is called
eval_modules = (model.rpn, model.backbone, model.roi_heads.box,)
fix_eval_modules(eval_modules)
# NOTE, we slow down the LR of the layers start with the names in slow_heads
if cfg.MODEL.ROI_RELATION_HEAD.PREDICTOR == "IMPPredictor":
slow_heads = ["roi_heads.relation.box_feature_extractor",
"roi_heads.relation.union_feature_extractor.feature_extractor",]
else:
slow_heads = []
# load pretrain layers to new layers
load_mapping = {"roi_heads.relation.box_feature_extractor" : "roi_heads.box.feature_extractor",
"roi_heads.relation.union_feature_extractor.feature_extractor" : "roi_heads.box.feature_extractor"}
if cfg.MODEL.ATTRIBUTE_ON:
load_mapping["roi_heads.relation.att_feature_extractor"] = "roi_heads.attribute.feature_extractor"
load_mapping["roi_heads.relation.union_feature_extractor.att_feature_extractor"] = "roi_heads.attribute.feature_extractor"
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
num_batch = cfg.SOLVER.IMS_PER_BATCH
optimizer = make_optimizer(cfg, model, logger, slow_heads=slow_heads, slow_ratio=10.0, rl_factor=float(num_batch))
scheduler = make_lr_scheduler(cfg, optimizer, logger)
debug_print(logger, 'end optimizer and shcedule')
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model, optimizer, opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
find_unused_parameters=True,
)
debug_print(logger, 'end distributed')
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk, custom_scheduler=True
)
# if there is certain checkpoint in output_dir, load it, else load pretrained detector
if checkpointer.has_checkpoint():
extra_checkpoint_data = checkpointer.load(cfg.MODEL.PRETRAINED_DETECTOR_CKPT,
update_schedule=cfg.SOLVER.UPDATE_SCHEDULE_DURING_LOAD)
arguments.update(extra_checkpoint_data)
if cfg.SOLVER.UPDATE_SCHEDULE_DURING_LOAD:
checkpointer.scheduler.last_epoch = extra_checkpoint_data["iteration"]
logger.info("update last epoch of scheduler to iter: {}".format(str(extra_checkpoint_data["iteration"])))
else:
# load_mapping is only used when we init current model from detection model.
checkpointer.load(cfg.MODEL.PRETRAINED_DETECTOR_CKPT, with_optim=False, load_mapping=load_mapping)
debug_print(logger, 'end load checkpointer')
train_data_loader = make_data_loader(
cfg,
mode='train',
is_distributed=distributed,
start_iter=arguments["iteration"],
)
val_data_loaders = make_data_loader(
cfg,
mode='val',
is_distributed=distributed,
)
debug_print(logger, 'end dataloader')
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
if cfg.SOLVER.PRE_VAL:
logger.info("Validate before training")
run_val(cfg, model, val_data_loaders, distributed, logger)
logger.info("Start training")
meters = MetricLogger(delimiter=" ")
max_iter = len(train_data_loader)
start_iter = arguments["iteration"]
start_training_time = time.time()
end = time.time()
print_first_grad = True
for iteration, (images, targets, _) in enumerate(train_data_loader, start_iter):
if any(len(target) < 1 for target in targets):
logger.error(f"Iteration={iteration + 1} || Image Ids used for training {_} || targets Length={[len(target) for target in targets]}" )
data_time = time.time() - end
iteration = iteration + 1
arguments["iteration"] = iteration
model.train()
fix_eval_modules(eval_modules)
images = images.to(device)
targets = [target.to(device) for target in targets]
loss_dict = model(images, targets)
losses = sum(loss for loss in loss_dict.values())
# reduce losses over all GPUs for logging purposes
loss_dict_reduced = reduce_loss_dict(loss_dict)
losses_reduced = sum(loss for loss in loss_dict_reduced.values())
meters.update(loss=losses_reduced, **loss_dict_reduced)
optimizer.zero_grad()
# Note: If mixed precision is not used, this ends up doing nothing
# Otherwise apply loss scaling for mixed-precision recipe
with amp.scale_loss(losses, optimizer) as scaled_losses:
scaled_losses.backward()
# add clip_grad_norm from MOTIFS, tracking gradient, used for debug
verbose = (iteration % cfg.SOLVER.PRINT_GRAD_FREQ) == 0 or print_first_grad # print grad or not
print_first_grad = False
clip_grad_norm([(n, p) for n, p in model.named_parameters() if p.requires_grad], max_norm=cfg.SOLVER.GRAD_NORM_CLIP, logger=logger, verbose=verbose, clip=True)
optimizer.step()
batch_time = time.time() - end
end = time.time()
meters.update(time=batch_time, data=data_time)
eta_seconds = meters.time.global_avg * (max_iter - iteration)
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
if iteration % 200 == 0 or iteration == max_iter:
logger.info(
meters.delimiter.join(
[
"eta: {eta}",
"iter: {iter}",
"{meters}",
"lr: {lr:.6f}",
"max mem: {memory:.0f}",
]
).format(
eta=eta_string,
iter=iteration,
meters=str(meters),
lr=optimizer.param_groups[-1]["lr"],
memory=torch.cuda.max_memory_allocated() / 1024.0 / 1024.0,
)
)
if iteration % checkpoint_period == 0:
checkpointer.save("model_{:07d}".format(iteration), **arguments)
if iteration == max_iter:
checkpointer.save("model_final", **arguments)
val_result = None # used for scheduler updating
if cfg.SOLVER.TO_VAL and iteration % cfg.SOLVER.VAL_PERIOD == 0:
logger.info("Start validating")
val_result = run_val(cfg, model, val_data_loaders, distributed, logger)
logger.info("Validation Result: %.4f" % val_result)
# scheduler should be called after optimizer.step() in pytorch>=1.1.0
# https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
if cfg.SOLVER.SCHEDULE.TYPE == "WarmupReduceLROnPlateau":
scheduler.step(val_result, epoch=iteration)
if scheduler.stage_count >= cfg.SOLVER.SCHEDULE.MAX_DECAY_STEP:
logger.info("Trigger MAX_DECAY_STEP at iteration {}.".format(iteration))
break
else:
scheduler.step()
total_training_time = time.time() - start_training_time
total_time_str = str(datetime.timedelta(seconds=total_training_time))
logger.info(
"Total training time: {} ({:.4f} s / it)".format(
total_time_str, total_training_time / (max_iter)
)
)
return model
def train(cfg, local_rank, distributed):
# Model logging
print_mlperf(key=mlperf_log.INPUT_BATCH_SIZE,
value=cfg.SOLVER.IMS_PER_BATCH)
print_mlperf(key=mlperf_log.BATCH_SIZE_TEST, value=cfg.TEST.IMS_PER_BATCH)
print_mlperf(key=mlperf_log.INPUT_MEAN_SUBTRACTION,
value=cfg.INPUT.PIXEL_MEAN)
print_mlperf(key=mlperf_log.INPUT_NORMALIZATION_STD,
value=cfg.INPUT.PIXEL_STD)
print_mlperf(key=mlperf_log.INPUT_RESIZE)
print_mlperf(key=mlperf_log.INPUT_RESIZE_ASPECT_PRESERVING)
print_mlperf(key=mlperf_log.MIN_IMAGE_SIZE, value=cfg.INPUT.MIN_SIZE_TRAIN)
print_mlperf(key=mlperf_log.MAX_IMAGE_SIZE, value=cfg.INPUT.MAX_SIZE_TRAIN)
print_mlperf(key=mlperf_log.INPUT_RANDOM_FLIP)
print_mlperf(key=mlperf_log.RANDOM_FLIP_PROBABILITY, value=0.5)
print_mlperf(key=mlperf_log.FG_IOU_THRESHOLD,
value=cfg.MODEL.RPN.FG_IOU_THRESHOLD)
print_mlperf(key=mlperf_log.BG_IOU_THRESHOLD,
value=cfg.MODEL.RPN.BG_IOU_THRESHOLD)
print_mlperf(key=mlperf_log.RPN_PRE_NMS_TOP_N_TRAIN,
value=cfg.MODEL.RPN.PRE_NMS_TOP_N_TRAIN)
print_mlperf(key=mlperf_log.RPN_PRE_NMS_TOP_N_TEST,
value=cfg.MODEL.RPN.PRE_NMS_TOP_N_TEST)
print_mlperf(key=mlperf_log.RPN_POST_NMS_TOP_N_TRAIN,
value=cfg.MODEL.RPN.FPN_POST_NMS_TOP_N_TRAIN)
print_mlperf(key=mlperf_log.RPN_POST_NMS_TOP_N_TEST,
value=cfg.MODEL.RPN.FPN_POST_NMS_TOP_N_TEST)
print_mlperf(key=mlperf_log.ASPECT_RATIOS,
value=cfg.MODEL.RPN.ASPECT_RATIOS)
print_mlperf(key=mlperf_log.BACKBONE, value=cfg.MODEL.BACKBONE.CONV_BODY)
print_mlperf(key=mlperf_log.NMS_THRESHOLD, value=cfg.MODEL.RPN.NMS_THRESH)
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
# Optimizer logging
print_mlperf(key=mlperf_log.OPT_NAME, value=mlperf_log.SGD_WITH_MOMENTUM)
print_mlperf(key=mlperf_log.OPT_LR, value=cfg.SOLVER.BASE_LR)
print_mlperf(key=mlperf_log.OPT_MOMENTUM, value=cfg.SOLVER.MOMENTUM)
print_mlperf(key=mlperf_log.OPT_WEIGHT_DECAY,
value=cfg.SOLVER.WEIGHT_DECAY)
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
arguments["save_checkpoints"] = cfg.SAVE_CHECKPOINTS
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader, iters_per_epoch = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
# set the callback function to evaluate and potentially
# early exit each epoch
if cfg.PER_EPOCH_EVAL:
per_iter_callback_fn = functools.partial(
mlperf_test_early_exit,
iters_per_epoch=iters_per_epoch,
tester=functools.partial(test, cfg=cfg),
model=model,
distributed=distributed,
min_bbox_map=cfg.MLPERF.MIN_BBOX_MAP,
min_segm_map=cfg.MLPERF.MIN_SEGM_MAP)
else:
per_iter_callback_fn = None
start_train_time = time.time()
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
per_iter_start_callback_fn=functools.partial(
mlperf_log_epoch_start, iters_per_epoch=iters_per_epoch),
per_iter_end_callback_fn=per_iter_callback_fn,
)
end_train_time = time.time()
total_training_time = end_train_time - start_train_time
print("&&&& MLPERF METRIC THROUGHPUT per GPU={:.4f} iterations / s".format(
(arguments["iteration"] * 1.0) / total_training_time))
return model
def train(cfg, local_rank, distributed):
model,head = build_sharedFC_face_trainer(cfg,local_rank)
device = torch.device(cfg.MODEL.DEVICE)
if cfg.MODEL.USE_SYNCBN:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
head_optimizer = make_optimizer(cfg, head)
head_scheduler = make_lr_scheduler(cfg, head_optimizer)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model, optimizer, opt_level=amp_opt_level)
head, head_optimizer = amp.initialize(head, head_optimizer, opt_level=amp_opt_level)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
head_checkpointer = DetectronCheckpointer(
cfg, head, head_optimizer, head_scheduler, output_dir, save_to_disk
)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
#### init transforms #####
transforms = T.Compose(
[
T.RandomCrop( (cfg.INPUT.SIZE_TRAIN[0], cfg.INPUT.SIZE_TRAIN[1]) ),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=cfg.INPUT.RGB_MEAN, std=cfg.INPUT.RGB_STD),
]
)
data_loader = make_face_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
transforms=transforms,
)
test_period = cfg.SOLVER.TEST_PERIOD
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
divs_nums = cfg.SOLVER.DIVS_NUMS_PER_BATCH
do_face_train_dist(
cfg,
[model,head],
data_loader,
None,
[optimizer,head_optimizer],
[scheduler,head_scheduler],
[checkpointer,head_checkpointer],
device,
checkpoint_period,
test_period,
arguments,
divs_nums,
)
return model
def train(cfg, local_rank, distributed, use_tensorboard=False):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model, optimizer, opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
test_period = cfg.SOLVER.TEST_PERIOD
if test_period > 0:
data_loader_val = make_data_loader(cfg, is_train=False, is_distributed=distributed, is_for_period=True)
else:
data_loader_val = None
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
if use_tensorboard:
meters = TensorboardLogger(
log_dir=cfg.TENSORBOARD_EXPERIMENT,
stage = 'train',
start_iter=arguments['iteration'],
delimiter=" ")
meters_val = TensorboardLogger(
log_dir=cfg.TENSORBOARD_EXPERIMENT,
stage = 'val',
start_iter=arguments['iteration'],
delimiter=" ")
else:
meters = MetricLogger(delimiter=" ")
meters_val = MetricLogger(delimiter=" ")
do_train(
cfg,
model,
data_loader,
data_loader_val,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
test_period,
arguments,
meters,
meters_val,
)
return model
def train(cfg, local_rank, distributed, logger):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg,
model,
logger,
rl_factor=float(cfg.SOLVER.IMS_PER_BATCH))
scheduler = make_lr_scheduler(cfg, optimizer)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
extra_checkpoint_data = checkpointer.load(
cfg.MODEL.WEIGHT,
update_schedule=cfg.SOLVER.UPDATE_SCHEDULE_DURING_LOAD)
arguments.update(extra_checkpoint_data)
train_data_loader = make_data_loader(
cfg,
mode='train',
is_distributed=distributed,
start_iter=arguments["iteration"],
)
val_data_loaders = make_data_loader(
cfg,
mode='val',
is_distributed=distributed,
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
if cfg.SOLVER.PRE_VAL:
logger.info("Validate before training")
run_val(cfg, model, val_data_loaders, distributed)
logger.info("Start training")
meters = MetricLogger(delimiter=" ")
max_iter = len(train_data_loader)
start_iter = arguments["iteration"]
start_training_time = time.time()
end = time.time()
for iteration, (images, targets, _) in enumerate(train_data_loader,
start_iter):
model.train()
if any(len(target) < 1 for target in targets):
logger.error(
f"Iteration={iteration + 1} || Image Ids used for training {_} || targets Length={[len(target) for target in targets]}"
)
data_time = time.time() - end
iteration = iteration + 1
arguments["iteration"] = iteration
scheduler.step()
images = images.to(device)
targets = [target.to(device) for target in targets]
loss_dict = model(images, targets)
losses = sum(loss for loss in loss_dict.values())
# reduce losses over all GPUs for logging purposes
loss_dict_reduced = reduce_loss_dict(loss_dict)
losses_reduced = sum(loss for loss in loss_dict_reduced.values())
meters.update(loss=losses_reduced, **loss_dict_reduced)
optimizer.zero_grad()
# Note: If mixed precision is not used, this ends up doing nothing
# Otherwise apply loss scaling for mixed-precision recipe
with amp.scale_loss(losses, optimizer) as scaled_losses:
scaled_losses.backward()
optimizer.step()
batch_time = time.time() - end
end = time.time()
meters.update(time=batch_time, data=data_time)
eta_seconds = meters.time.global_avg * (max_iter - iteration)
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
if iteration % 200 == 0 or iteration == max_iter:
logger.info(
meters.delimiter.join([
"eta: {eta}",
"iter: {iter}",
"{meters}",
"lr: {lr:.6f}",
"max mem: {memory:.0f}",
]).format(
eta=eta_string,
iter=iteration,
meters=str(meters),
lr=optimizer.param_groups[0]["lr"],
memory=torch.cuda.max_memory_allocated() / 1024.0 / 1024.0,
))
if cfg.SOLVER.TO_VAL and iteration % cfg.SOLVER.VAL_PERIOD == 0:
logger.info("Start validating")
run_val(cfg, model, val_data_loaders, distributed)
if iteration % checkpoint_period == 0:
checkpointer.save("model_{:07d}".format(iteration), **arguments)
if iteration == max_iter:
checkpointer.save("model_final", **arguments)
total_training_time = time.time() - start_training_time
total_time_str = str(datetime.timedelta(seconds=total_training_time))
logger.info("Total training time: {} ({:.4f} s / it)".format(
total_time_str, total_training_time / (max_iter)))
return model
def train(cfg, args, DatasetCatalog=None):
if len(cfg.DATASETS.TRAIN) == 0 or not args.train:
return None
local_rank = args.local_rank
distributed = args.distributed
model = build_detection_model(cfg)
# for key, value in model.named_parameters():
# print(key, value.requires_grad)
if hasattr(args, 'train_last_layer'):
if args.train_last_layer:
listofkeys = [
'cls_score.bias', 'cls_score.weight', 'bbox_pred.bias',
'bbox_pred.weight', 'mask_fcn_logits.bias',
'mask_fcn_logits.weight'
]
for key, value in model.named_parameters():
value.requires_grad = False
for k in listofkeys:
if k in key:
value.requires_grad = True
# for key, value in model.named_parameters():
# print(key, value.requires_grad)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# Initialize mixed-precision training
if cfg.MODEL.DEVICE == 'cuda':
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
cfg.OUTPUT_DIR, save_to_disk)
extra_checkpoint_data = checkpointer.load(
cfg.MODEL.WEIGHT,
force_load_external_checkpoint=False,
copy_weight_from_head_box=args.copy_weight_from_head_box)
arguments = {}
arguments.update(extra_checkpoint_data)
data_loader = make_data_loader(
cfg,
args,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
DatasetCatalog=DatasetCatalog,
)
if cfg.SOLVER.TEST_PERIOD > 0:
data_loader_val = make_data_loader(
cfg,
args,
is_train=False,
is_distributed=distributed,
is_for_period=True,
start_iter=arguments["iteration"],
DatasetCatalog=DatasetCatalog,
)
else:
data_loader_val = None
do_train(
model,
cfg,
data_loader,
data_loader_val,
optimizer,
scheduler,
checkpointer,
device,
cfg.SOLVER.CHECKPOINT_PERIOD,
cfg.SOLVER.TEST_PERIOD,
arguments,
cfg.OUTPUT_DIR,
args.visualize_loss,
args.vis_title,
args.iters_per_epoch,
)
return model
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, None, None, output_dir,
save_to_disk)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
# arguments.update(extra_checkpoint_data)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
logger = logging.getLogger("maskrcnn_benchmark.trainer")
if cfg.MODEL.META_ARCHITECTURE == 'AdaptionRCNN':
logger.info('AdaptionRCNN trainer is adapted!')
cross_do_train(
cfg,
model,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
distributed,
)
elif cfg.MODEL.META_ARCHITECTURE == 'GeneralizedRCNN':
logger.info('GeneralizedRCNN trainer is adapted!')
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
do_train(
cfg,
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
distributed,
)
return model
def train(cfg, local_rank, distributed):
model = build_face_trainer(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
if cfg.MODEL.USE_SYNCBN:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
#### init transforms #####
transforms = T.Compose([
T.RandomCrop((cfg.INPUT.SIZE_TRAIN[0], cfg.INPUT.SIZE_TRAIN[1])),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=cfg.INPUT.RGB_MEAN, std=cfg.INPUT.RGB_STD),
])
data_loader = make_face_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
transforms=transforms,
)
test_period = cfg.SOLVER.TEST_PERIOD
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
divs_nums = cfg.SOLVER.DIVS_NUMS_PER_BATCH
BUILD_FACE_TRAINER(cfg)(
cfg,
model,
data_loader,
None,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
test_period,
arguments,
divs_nums,
)
return model
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
# Initialize mixed-precision training
use_mixed_precision = cfg.DTYPE == "float16"
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
arguments = {}
arguments["iteration"] = 0
arguments['phase'] = 1
arguments['plot_median'], arguments['plot_global_avg'] = defaultdict(
list), defaultdict(list)
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler,
output_dir, save_to_disk)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
test_period = cfg.SOLVER.TEST_PERIOD
if test_period > 0:
data_loader_val = make_data_loader(cfg,
is_train=False,
is_distributed=distributed,
is_for_period=True)
else:
data_loader_val = None
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
if arguments['phase'] == 1:
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
phase=1,
)
do_train(
cfg,
model,
data_loader,
data_loader_val,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
test_period,
arguments,
training_phase=1,
)
arguments["iteration"] = 0
arguments["phase"] = 2
data_loader_phase2 = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
phase=2,
)
do_train(
cfg,
model,
data_loader_phase2,
data_loader_val,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
test_period,
arguments,
training_phase=2,
)
return model
