def __init__(
self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
):
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
# self.model = torch.load('/home/zoey/nas/zoey/github/maskrcnn-benchmark/tinycoco/model_0010000.pth')
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.min_image_size = min_image_size
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
_ = checkpointer.load(
'/home/zoey/nas/zoey/github/maskrcnn-benchmark/tinycoco/model_0010000.pth'
)
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def test(cfg, model, distributed):
output_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
if distributed:
model = model.module
torch.cuda.empty_cache() # TODO check if it helps
iou_types = ("bbox", )
if cfg.MODEL.MASK_ON:
iou_types = iou_types + ("segm", )
output_folders = [None] * len(cfg.DATASETS.TEST)
if cfg.OUTPUT_DIR:
dataset_names = cfg.DATASETS.TEST
for idx, dataset_name in enumerate(dataset_names):
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference",
dataset_name)
mkdir(output_folder)
output_folders[idx] = output_folder
data_loaders_val = make_data_loader(cfg,
is_train=False,
is_distributed=distributed)
for output_folder, data_loader_val in zip(output_folders,
data_loaders_val):
inference(
model,
data_loader_val,
iou_types=iou_types,
box_only=cfg.MODEL.RPN_ONLY,
device=cfg.MODEL.DEVICE,
expected_results=cfg.TEST.EXPECTED_RESULTS,
expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL,
output_folder=output_folder,
maskiou_on=cfg.MODEL.MASKIOU_ON)
synchronize()
def __init__(self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224):
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
for n, m in self.model.named_modules():
if n == "roi_heads":
m.register_forward_hook(hook)
self.min_image_size = min_image_size
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def __init__(
self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
):
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.min_image_size = min_image_size
checkpointer = DetectronCheckpointer(cfg, self.model)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
# used to make colors for each class
self.palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def main():
parser = argparse.ArgumentParser(
description="PyTorch Object Detection Inference")
parser.add_argument(
"--config-file",
default=
"/home/qinjianbo/SRC/maskrcnn-benchmark/configs/e2e_faster_rcnn_R_50_FPN_1x.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument("--local_rank", type=int, default=0)
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
num_gpus = int(
os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
distributed = num_gpus > 1
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl",
init_method="env://")
synchronize()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
save_dir = ""
logger = setup_logger("maskrcnn_benchmark", save_dir, get_rank())
logger.info("Using {} GPUs".format(num_gpus))
logger.info(cfg)
logger.info("Collecting env info (might take some time)")
logger.info("\n" + collect_env_info())
model = build_detection_model(cfg)
model.to(cfg.MODEL.DEVICE)
output_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
output_folders = [None] * len(cfg.DATASETS.TEST)
'''
dataset_names = cfg.DATASETS.TEST
if cfg.OUTPUT_DIR:
for idx, dataset_name in enumerate(dataset_names):
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
mkdir(output_folder)
output_folders[idx] = output_folder
'''
# data_loaders_inference = make_data_loader(cfg, is_train=False, is_distributed=distributed)
return model, cfg, distributed
def train(cfg, local_rank, distributed, tb_logger):
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,
resume=cfg.SOLVER.RESUME)
if cfg.SOLVER.RESUME:
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,
tb_logger,
cfg,
local_rank,
)
return model
def __init__(
self,
cfg,
weights,
confidence_threshold=0.5,
min_image_size=224,
):
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.min_image_size = min_image_size
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
_ = checkpointer.load(weights)
self.transforms = self.build_transform()
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
def __init__(
self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
):
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.min_image_size = min_image_size
checkpointer = DetectronCheckpointer(cfg, self.model)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.checkpointer = checkpointer
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
self.CATEGORIES = COCO_CATEGORIES if cfg.DATASETS.TEST[
0][:4] == 'coco' else VOC_CATEGORIES
def train(cfg, local_rank, distributed):
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
print(model)
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 = os.path.join(cfg.OUTPUT_DIR, cfg.FILE)
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_train = 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,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
val_period = cfg.SOLVER.VAL_PERIOD
do_train(
model,
data_loader_train,
data_loader_val,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
val_period,
arguments,
distributed,
)
return model
def test(self, output_dir=None, model_to_test=None):
if output_dir is not None:
self.cfg.OUTPUT_DIR = output_dir
model = build_detection_model(self.cfg)
device = torch.device(self.cfg.MODEL.DEVICE)
model.to(device)
arguments = {}
arguments["iteration"] = 0
output_dir = self.cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
self.cfg, model, None, None, output_dir, save_to_disk
)
if model_to_test is not None:
self.cfg.MODEL.WEIGHT = model_to_test
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))
extra_checkpoint_data = checkpointer.load(model_path, use_latest=False)
checkpointer.optimizer = make_optimizer(self.cfg, checkpointer.model)
checkpointer.scheduler = make_lr_scheduler(self.cfg, checkpointer.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(checkpointer.model, checkpointer.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,
)
synchronize()
_ = inference( # The result can be used for additional logging, e. g. for TensorBoard
model,
# The method changes the segmentation mask format in a data loader,
# so every time a new data loader is created:
make_data_loader(self.cfg, is_train=False, is_distributed=(get_world_size() > 1), is_target_task=self.is_target_task),
dataset_name="[Test]",
iou_types=("bbox",),
box_only=False if cfg.MODEL.RETINANET_ON else cfg.MODEL.RPN_ONLY,
device=cfg.MODEL.DEVICE,
expected_results=cfg.TEST.EXPECTED_RESULTS,
expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL,
output_folder=None,
is_target_task=self.is_target_task,
)
synchronize()
logger = logging.getLogger("maskrcnn_benchmark")
logger.handlers=[]
def test(cfg, args, output_dir):
torch.cuda.empty_cache()
# Construct model graph
model = build_siammot(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
# Load model params
model_file = args.model_file
checkpointer = DetectronCheckpointer(cfg, model, save_dir=model_file)
if os.path.isfile(model_file):
_ = checkpointer.load(model_file)
elif os.path.isdir(model_file):
_ = checkpointer.load(use_latest=True)
else:
raise KeyError("No checkpoint is found")
# Load testing dataset
dataset_key = args.test_dataset
dataset, modality = load_dataset_anno(cfg, dataset_key, args.set)
dataset = sorted(dataset)
# do inference on dataset
data_filter_fn = build_data_filter_fn(dataset_key)
# load public detection
public_detection = None
if cfg.INFERENCE.USE_GIVEN_DETECTIONS:
public_detection = load_public_detection(cfg, dataset_key)
dataset_inference = DatasetInference(cfg, model, dataset, output_dir, data_filter_fn, public_detection)
dataset_inference()
def __init__(
self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
min_image_size=224,
):
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.min_image_size = min_image_size
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
def __init__(self, cfg, pretrained_model_path, build_transform=False):
super(ResnetEncoder, self).__init__()
# basic properties
self.cfg = cfg
self.transforms = None
if build_transform:
self.transforms = self.build_transform()
self.device = torch.device(cfg.MODEL.DEVICE)
# loading mask rcnn
self.maskrcnn = build_detection_model(cfg)
self.maskrcnn.eval()
device = torch.device(cfg.MODEL.DEVICE)
self.maskrcnn.to(device)
self.checkpointer = DetectronCheckpointer(cfg,
self.maskrcnn,
save_dir='.')
_ = self.checkpointer.load(pretrained_model_path)
# freeze gradients for mask rcnn
for param in self.maskrcnn.backbone.parameters():
param.requires_grad = False
for param in self.maskrcnn.rpn.parameters():
param.requires_grad = False
for param in self.maskrcnn.roi_heads.parameters():
param.requires_grad = False
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 __init__(
self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
):
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.min_image_size = min_image_size
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
logging.info('model loaded from: {}'.format(cfg.MODEL.WEIGHT))
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def train(cfg, local_rank, distributed, use_tensorboard=False, logger=None, start_iter=0):
arguments = {"iteration": start_iter}
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, start_iter)
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, start_iter)
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 __init__(
self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=800,
):
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.min_image_size = min_image_size
checkpointer = DetectronCheckpointer(cfg, self.model)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
# used to make colors for each class
self.palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
self.color_table = np.array([[162, 109, 35], [69, 94, 183], [72, 161, 198],
[82, 158, 127], [120, 72, 122], [105, 124, 135]]*10)
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def __init__(
self,
cfg,
# show_mask_heatmaps=False,
# masks_per_dim=2,
):
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.min_image_size = cfg.INPUT.MIN_SIZE_TEST
self.max_image_size = cfg.INPUT.MAX_SIZE_TEST
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
# mask_threshold = -1 if show_mask_heatmaps else 0.5
# self.masker = Masker(threshold=mask_threshold, padding=1)
# used to make colors for each class
self.palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
self.cpu_device = torch.device("cpu")
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, 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
if cfg.USE_TENSORBOARD_LOGS:
meters = TensorboardLogger(
log_dir=os.path.join(output_dir, 'tensorboard_logs'),
start_iter=arguments['iteration'],
delimiter=" ",
)
else:
meters = MetricLogger(delimiter=" ")
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
meters,
)
return model
def train(cfg, local_rank, distributed):
# 创建GeneralizedRCNN()对象
# detectors.py --> generalized_rcnn.py
model = build_detection_model(cfg)
# print(model)
# '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
# 获取输出的文件夹路径, 默认为 '.'
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) # 字典的update方法, 对字典的键值进行更新
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 main():
parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference")
parser.add_argument(
"--config-file",
default="/private/home/fmassa/github/detectron.pytorch_v2/configs/e2e_faster_rcnn_R_50_C4_1x_caffe2.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument("--local_rank", type=int, default=0)
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
distributed = num_gpus > 1
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(
backend="nccl", init_method="env://"
)
synchronize()
cfg.merge_from_list(args.opts)
output_dir = cfg.OUTPUT_DIR
config_file = os.path.join(output_dir, "runtime_config.yaml")
if args.config_file != "":
config_file = args.config_file
cfg.merge_from_file(config_file)
cfg.merge_from_list(args.opts)
adjustment_for_relation(cfg)
cfg.freeze()
save_dir = ""
logger = setup_logger("maskrcnn_benchmark", save_dir, get_rank())
logger.info("Using {} GPUs".format(num_gpus))
logger.info(cfg)
logger.info("Collecting env info (might take some time)")
logger.info("\n" + collect_env_info())
data_loader_val = make_data_loader(cfg, is_train=False, is_distributed=distributed)
model = build_detection_model(cfg)
model.to(cfg.MODEL.DEVICE)
checkpoint_output_dir = os.path.join(output_dir, 'checkpoints')
checkpointer = DetectronCheckpointer(cfg, model, save_dir=checkpoint_output_dir)
checkpoint, ckpt_fname = checkpointer.load(cfg.MODEL.WEIGHT)
results_dict = compute_on_dataset(model, data_loader_val[0], cfg.MODEL.DEVICE)
predictions = _accumulate_predictions_from_multiple_gpus(results_dict)
torch.save(predictions, '/p300/flickr30k_images/flickr30k_anno/precomp_proposals_nms1e5.pth')
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, 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"],
)
if checkpointer.classes is None:
for ds in data_loader.dataset.datasets:
ds.find_classes()
checkpointer.classes = data_loader.dataset.datasets[0].class_to_ind
else:
print("Loading classes from file")
print(checkpointer.classes)
for ds in data_loader.dataset.datasets:
ds.class_to_ind = checkpointer.classes
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):
# use following line to avoid shared file limit
# torch.multiprocessing.set_sharing_strategy('file_system')
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)
# Convert Model for SyncBN
if cfg.SYNCBN:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
# 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"],
)
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):
# ################################################################### fusion_factors # add by G
if cfg.MODEL.FPN.STATISTICS_ALPHA_ON == True:
sta_module = StaAlphaModule(cfg)
fusion_factors = sta_module.process()
else:
fusion_factors = cfg.MODEL.FPN.FUSION_FACTORS
# ################################################################### fusion_factors # add by G
model = build_detection_model(cfg, fusion_factors)
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, 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, 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)
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
if use_tensorboard:
arguments["tb_log_dir"] = cfg.TENSORBOARD_LOGDIR
arguments["tb_exp_name"] = cfg.TENSORBOARD_EXP_NAME
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"],
)
print(data_loader.dataset)
for iteration, (images, targets, _) in enumerate(data_loader, 0):
print(">>>>> train iteration:", iteration)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
)
return model
def _build_detector(self):
self.model = build_detection_model(self.cfg)
self.model.eval()
self.model.to(self.device)
checkpointer = DetectronCheckpointer(self.cfg, self.model)
_ = checkpointer.load(self.model_path)
self.transforms = self._build_transform()
def __init__(self):
super(MaskRCNN_Benchmark, self).__init__()
self.avgpool = nn.AdaptiveAvgPool2d(output_size=1)
self.model = build_detection_model(cfg)
# load the pre-trained model
checkpointer = DetectronCheckpointer(cfg, self.model)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
# make sure maskrcnn_benchmark is in eval mode
self.model.eval()
def build_and_load_model(self):
cfg = self.cfg
model = build_detection_model(cfg)
model.to(cfg.MODEL.DEVICE)
checkpointer = DetectronCheckpointer(cfg,
model,
save_dir=cfg.OUTPUT_DIR)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
model.eval()
self.model = 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, optimizer, scheduler,
output_dir, save_to_disk)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
arguments.update(extra_checkpoint_data)
#HACK: force the steps, could not change the lr from ckpt now.
scheduler.milestones = cfg.SOLVER.STEPS
# change lr
#lr_ratio = cfg.SOLVER.BASE_LR / scheduler.base_lrs[-1]
#scheduler.base_lrs = [ base_lr * lr_ratio for base_lr in self.base_lrs ]
data_loader = make_data_loader(
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
tbwriter = SummaryWriter(cfg.OUTPUT_DIR)
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
checkpoint_period,
arguments,
tbwriter,
)
return model
def __init__(self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
weight_loading=None):
self.cfg = cfg.clone()
# dynamically load labels.json in log directory
self.CATEGORIES = ["__background"]
if 'wolf' in self.cfg.DATASETS.TEST[0]:
with open('../log/wolf_labels.json') as f:
labels = json.load(f)
else:
with open('../log/coco_labels.json') as f:
labels = json.load(f)
for id in labels:
self.CATEGORIES.append(labels[id])
print(self.CATEGORIES)
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
print('self.device: {}'.format(self.device))
self.model.to(self.device)
self.min_image_size = min_image_size
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
if weight_loading:
print('Loading weight from {}.'.format(weight_loading))
_ = checkpointer._load_model(torch.load(weight_loading))
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
self.cpu_device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
# self.cpu_device = torch.device("cpu")
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1,
2**21 - 1]).to(self.cpu_device)
self.confidence_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def load_model(cfg, cuda=True):
device = torch.device("cuda" if cuda else "cpu")
cfg = cfg.clone()
model = build_detection_model(cfg)
model.eval()
model.to(device)
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, model, save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
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.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 main():
parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference")
parser.add_argument(
"--config-file",
default="/private/home/fmassa/github/detectron.pytorch_v2/configs/e2e_faster_rcnn_R_50_C4_1x_caffe2.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument("--local_rank", type=int, default=0)
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
distributed = num_gpus > 1
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.deprecated.init_process_group(
backend="nccl", init_method="env://"
)
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
save_dir = ""
logger = setup_logger("maskrcnn_benchmark", save_dir, get_rank())
logger.info("Using {} GPUs".format(num_gpus))
logger.info(cfg)
logger.info("Collecting env info (might take some time)")
logger.info("\n" + collect_env_info())
model = build_detection_model(cfg)
model.to(cfg.MODEL.DEVICE)
output_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
iou_types = ("bbox",)
if cfg.MODEL.MASK_ON:
iou_types = iou_types + ("segm",)
output_folders = [None] * len(cfg.DATASETS.TEST)
dataset_names = cfg.DATASETS.TEST
if cfg.OUTPUT_DIR:
for idx, dataset_name in enumerate(dataset_names):
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
mkdir(output_folder)
output_folders[idx] = output_folder
data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed)
for output_folder, dataset_name, data_loader_val in zip(output_folders, dataset_names, data_loaders_val):
inference(
model,
data_loader_val,
dataset_name=dataset_name,
iou_types=iou_types,
box_only=cfg.MODEL.RPN_ONLY,
device=cfg.MODEL.DEVICE,
expected_results=cfg.TEST.EXPECTED_RESULTS,
expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL,
output_folder=output_folder,
)
synchronize()