def setup_custom_environment(custom_module_path):
"""Load custom environment setup from a Python source file and run the setup
function.
"""
module = import_file("maskrcnn_benchmark.utils.env.custom_module", custom_module_path)
assert hasattr(module, "setup_environment") and callable(
module.setup_environment
), (
"Custom environment module defined in {} does not have the "
"required callable attribute 'setup_environment'."
).format(
custom_module_path
)
module.setup_environment()
def _load_file(self, f):
# catalog lookup
if f.startswith("catalog://"):
paths_catalog = import_file(
"maskrcnn_benchmark.config.paths_catalog", self.cfg.PATHS_CATALOG, True
)
catalog_f = paths_catalog.ModelCatalog.get(f[len("catalog://") :])
self.logger.info("{} points to {}".format(f, catalog_f))
f = catalog_f
# download url files
if f.startswith("http"):
# if the file is a url path, download it and cache it
cached_f = cache_url(f)
self.logger.info("url {} cached in {}".format(f, cached_f))
f = cached_f
# convert Caffe2 checkpoint from pkl
if f.endswith(".pkl"):
return load_c2_format(self.cfg, f)
# load native detectron.pytorch checkpoint
loaded = super(DetectronCheckpointer, self)._load_file(f)
if "model" not in loaded:
loaded = dict(model=loaded)
return loaded
def make_data_loader(cfg, is_train=True, is_distributed=False, start_iter=0):
num_gpus = get_world_size()
if is_train:
images_per_batch = cfg.SOLVER.IMS_PER_BATCH
assert (
images_per_batch % num_gpus == 0
), "SOLVER.IMS_PER_BATCH ({}) must be divisible by the number "
"of GPUs ({}) used.".format(images_per_batch, num_gpus)
images_per_gpu = images_per_batch // num_gpus
shuffle = True
num_iters = cfg.SOLVER.MAX_ITER
else:
images_per_batch = cfg.TEST.IMS_PER_BATCH
assert (
images_per_batch % num_gpus == 0
), "TEST.IMS_PER_BATCH ({}) must be divisible by the number "
"of GPUs ({}) used.".format(images_per_batch, num_gpus)
images_per_gpu = images_per_batch // num_gpus
shuffle = False if not is_distributed else True
num_iters = None
start_iter = 0
if images_per_gpu > 1:
logger = logging.getLogger(__name__)
logger.warning(
"When using more than one image per GPU you may encounter "
"an out-of-memory (OOM) error if your GPU does not have "
"sufficient memory. If this happens, you can reduce "
"SOLVER.IMS_PER_BATCH (for training) or "
"TEST.IMS_PER_BATCH (for inference). For training, you must "
"also adjust the learning rate and schedule length according "
"to the linear scaling rule. See for example: "
"https://github.com/facebookresearch/Detectron/blob/master/configs/getting_started/tutorial_1gpu_e2e_faster_rcnn_R-50-FPN.yaml#L14"
)
# group images which have similar aspect ratio. In this case, we only
# group in two cases: those with width / height > 1, and the other way around,
# but the code supports more general grouping strategy
aspect_grouping = [1] if cfg.DATALOADER.ASPECT_RATIO_GROUPING else []
paths_catalog = import_file(
"maskrcnn_benchmark.config.paths_catalog", cfg.PATHS_CATALOG, True
)
DatasetCatalog = paths_catalog.DatasetCatalog
dataset_list = cfg.DATASETS.TRAIN if is_train else cfg.DATASETS.TEST
transforms = build_transforms(cfg, is_train)
datasets = build_dataset(dataset_list, transforms, DatasetCatalog, is_train)
data_loaders = []
for dataset in datasets:
sampler = make_data_sampler(dataset, shuffle, is_distributed)
batch_sampler = make_batch_data_sampler(
dataset, sampler, aspect_grouping, images_per_gpu, num_iters, start_iter
)
collator = BatchCollator(cfg.DATALOADER.SIZE_DIVISIBILITY)
num_workers = cfg.DATALOADER.NUM_WORKERS
data_loader = torch.utils.data.DataLoader(
dataset,
num_workers=num_workers,
batch_sampler=batch_sampler,
collate_fn=collator,
)
data_loaders.append(data_loader)
if is_train:
# during training, a single (possibly concatenated) data_loader is returned
assert len(data_loaders) == 1
return data_loaders[0]
return data_loaders
def make_data_loader(cfg, is_train=True, is_distributed=False, start_iter=0):
num_gpus = get_world_size()
if is_train:
images_per_batch = cfg.SOLVER.IMS_PER_BATCH
assert (
images_per_batch % num_gpus == 0
), "SOLVER.IMS_PER_BATCH ({}) must be divisible by the number of GPUs ({}) used.".format(
images_per_batch, num_gpus)
images_per_gpu = images_per_batch // num_gpus
shuffle = True
num_iters = cfg.SOLVER.MAX_ITER
else:
images_per_batch = cfg.TEST.IMS_PER_BATCH
assert (
images_per_batch % num_gpus == 0
), "TEST.IMS_PER_BATCH ({}) must be divisible by the number of GPUs ({}) used.".format(
images_per_batch, num_gpus)
images_per_gpu = images_per_batch // num_gpus
shuffle = False if not is_distributed else True
num_iters = None
start_iter = 0
if images_per_gpu > 1:
logger = logging.getLogger(__name__)
logger.warning(
"When using more than one image per GPU you may encounter "
"an out-of-memory (OOM) error if your GPU does not have "
"sufficient memory. If this happens, you can reduce "
"SOLVER.IMS_PER_BATCH (for training) or "
"TEST.IMS_PER_BATCH (for inference). For training, you must "
"also adjust the learning rate and schedule length according "
"to the linear scaling rule. See for example: "
"https://github.com/facebookresearch/Detectron/blob/master/configs/getting_started/tutorial_1gpu_e2e_faster_rcnn_R-50-FPN.yaml#L14"
)
# group images which have similar aspect ratio. In this case, we only
# group in two cases: those with width / height > 1, and the other way around,
# but the code supports more general grouping strategy
aspect_grouping = [1] if cfg.DATALOADER.ASPECT_RATIO_GROUPING else []
paths_catalog = import_file("maskrcnn_benchmark.config.paths_catalog",
cfg.PATHS_CATALOG, True)
DatasetCatalog = paths_catalog.DatasetCatalog
dataset_list = cfg.DATASETS.TRAIN if is_train else cfg.DATASETS.TEST
# If bbox aug is enabled in testing, simply set transforms to None and we will apply transforms later
transforms = None if not is_train and cfg.TEST.BBOX_AUG.ENABLED else build_transforms(
cfg, is_train)
datasets = build_dataset(dataset_list, transforms, DatasetCatalog,
is_train)
if is_train:
# save category_id to label name mapping
save_labels(datasets, cfg.OUTPUT_DIR)
data_loaders = []
for dataset in datasets:
sampler = make_data_sampler(dataset, shuffle, is_distributed)
batch_sampler = make_batch_data_sampler(dataset, sampler,
aspect_grouping,
images_per_gpu, num_iters,
start_iter)
collator = BBoxAugCollator() if not is_train and cfg.TEST.BBOX_AUG.ENABLED else \
BatchCollator(cfg.DATALOADER.SIZE_DIVISIBILITY)
num_workers = cfg.DATALOADER.NUM_WORKERS
data_loader = torch.utils.data.DataLoader(
dataset,
num_workers=num_workers,
batch_sampler=batch_sampler,
collate_fn=collator,
)
data_loaders.append(data_loader)
if is_train:
# during training, a single (possibly concatenated) data_loader is returned
assert len(data_loaders) == 1
return data_loaders[0]
return data_loaders
def make_data_loader(cfg, is_train=True, is_distributed=False, start_iter=0):
num_gpus = get_world_size()
if is_train:
images_per_batch = cfg.SOLVER.IMS_PER_BATCH # 2
assert (images_per_batch % num_gpus == 0
), "SOLVER.IMS_PER_BATCH ({}) must be divisible by the number "
"of GPUs ({}) used.".format(images_per_batch, num_gpus)
images_per_gpu = images_per_batch // num_gpus
shuffle = True
num_iters = cfg.SOLVER.MAX_ITER # 720000
else:
images_per_batch = cfg.TEST.IMS_PER_BATCH # 1
assert (images_per_batch % num_gpus == 0
), "TEST.IMS_PER_BATCH ({}) must be divisible by the number "
"of GPUs ({}) used.".format(images_per_batch, num_gpus)
images_per_gpu = images_per_batch // num_gpus
shuffle = False if not is_distributed else True
num_iters = None
start_iter = 0
if images_per_gpu > 1:
logger = logging.getLogger(__name__)
logger.warning(
"When using more than one image per GPU you may encounter "
"an out-of-memory (OOM) error if your GPU does not have "
"sufficient memory. If this happens, you can reduce "
"SOLVER.IMS_PER_BATCH (for training) or "
"TEST.IMS_PER_BATCH (for inference). For training, you must "
"also adjust the learning rate and schedule length according "
"to the linear scaling rule. See for example: "
"https://github.com/facebookresearch/Detectron/blob/master/configs/getting_started/tutorial_1gpu_e2e_faster_rcnn_R-50-FPN.yaml#L14"
)
# group images which have similar aspect ratio. In this case, we only
# group in two cases: those with width / height > 1, and the other way around,
# but the code supports more general grouping strategy
# 默认为True
aspect_grouping = [1] if cfg.DATALOADER.ASPECT_RATIO_GROUPING else []
paths_catalog = import_file("maskrcnn_benchmark.config.paths_catalog",
cfg.PATHS_CATALOG, True)
# maskrcnn_benchmark.config.paths_catalog.py
DatasetCatalog = paths_catalog.DatasetCatalog # 将DatasetCatalog对象赋值给该变量
# ("coco_2014_train", "coco_2014_valminusminival") for train
# ("coco_2014_minival",) for test
dataset_list = cfg.DATASETS.TRAIN if is_train else cfg.DATASETS.TEST
# transform中传入的target指的是image对应的BoxList对象
transforms = build_transforms(cfg, is_train)
# 创建COCODataset对象
datasets = build_dataset(dataset_list, transforms, DatasetCatalog,
is_train)
data_loaders = []
for dataset in datasets:
# 这里创建RandomSampler对象
sampler = make_data_sampler(dataset, shuffle, is_distributed)
# 创建BatchSampler对象
batch_sampler = make_batch_data_sampler(dataset, sampler,
aspect_grouping,
images_per_gpu, num_iters,
start_iter)
# todo: 目前不清楚这个干啥的, 看看再说
collator = BatchCollator(cfg.DATALOADER.SIZE_DIVISIBILITY) # 32
# Number of data loading threads, 4
num_workers = cfg.DATALOADER.NUM_WORKERS
data_loader = torch.utils.data.DataLoader(
dataset,
num_workers=num_workers,
batch_sampler=batch_sampler,
collate_fn=collator,
)
data_loaders.append(data_loader)
if is_train:
# during training, a single (possibly concatenated) data_loader is returned
assert len(data_loaders) == 1
return data_loaders[0]
return data_loaders
def inference(
model,
data_loader,
dataset_name,
iou_types=("bbox", ),
box_only=False,
device=torch.device("cuda"),
expected_results=0,
expected_results_sigma_tol=0,
output_folder=None,
cfg=None,
bbox_aug=False,
visualize_results=False,
visualization_label="coco",
only_visualization=False,
):
num_devices = get_world_size()
logger = logging.getLogger("maskrcnn_benchmark.inference")
dataset = data_loader.dataset
logger.info("Start evaluation on {} dataset({} images).".format(
dataset_name, len(dataset)))
total_timer = Timer()
inference_timer = Timer()
total_timer.tic()
roi_predictions, img_predictions, attention_maps = compute_on_dataset(
model, data_loader, device, bbox_aug=bbox_aug, timer=inference_timer)
# wait for all processes to complete before measuring the time
synchronize()
total_time = total_timer.toc()
total_time_str = get_time_str(total_time)
logger.info(
"Total run time: {} ({} s / img per device, on {} devices)".format(
total_time_str, total_time * num_devices / len(dataset),
num_devices))
total_infer_time = get_time_str(inference_timer.total_time)
logger.info(
"Model inference time: {} ({} s / img per device, on {} devices)".
format(
total_infer_time,
inference_timer.total_time * num_devices / len(dataset),
num_devices,
))
if roi_predictions:
roi_predictions = _accumulate_predictions_from_multiple_gpus(
roi_predictions)
if img_predictions:
img_predictions = _accumulate_predictions_from_multiple_gpus(
img_predictions)
if attention_maps:
attention_maps = _accumulate_predictions_from_multiple_gpus(
attention_maps)
if not is_main_process():
return
if roi_predictions and len(roi_predictions) > 0:
for prediction in roi_predictions:
if prediction.has_field("pred_scores"):
prediction.add_field('second_scores',
prediction.get_field('pred_scores'))
del prediction.extra_fields["pred_scores"]
if prediction.has_field("pred_labels"):
prediction.add_field('second_labels',
prediction.get_field('pred_labels'))
del prediction.extra_fields["pred_labels"]
if output_folder:
torch.save(roi_predictions,
os.path.join(output_folder, "roi_predictions.pth"))
print('Visualize results')
if output_folder and visualize_results:
categories = import_file(
"maskrcnn_benchmark.data.datasets.categories.{}_categories".
format(visualization_label),
os.path.join(
os.path.dirname(os.path.dirname(cfg.PATHS_CATALOG)),
'data', 'categories',
'{}_categories.py'.format(visualization_label)), True)
visualizer = Visualizer(categories=categories.CATEGORIES, cfg=cfg)
visualizer.visualize_attentions(
attention_maps, dataset,
os.path.join(output_folder, 'attention_map'))
visualizer.visualize_predictions(
roi_predictions, dataset,
os.path.join(output_folder, 'visualization'))
if only_visualization:
return
extra_args = dict(
box_only=box_only,
iou_types=iou_types,
expected_results=expected_results,
expected_results_sigma_tol=expected_results_sigma_tol,
)
print('ROI: Evaluate')
evaluate_roi(dataset=dataset,
predictions=roi_predictions,
output_folder=output_folder,
**extra_args)
if img_predictions and len(img_predictions) > 0:
if output_folder:
torch.save(img_predictions,
os.path.join(output_folder, "img_predictions.pth"))
print('IMAGE: Evaluate')
evaluate_img(dataset=dataset,
predictions=img_predictions,
output_folder=output_folder)
def main(args):
num_gpus = get_num_gpus()
args.config_file = os.path.join(
info['training_dir'], 'e2e_faster_rcnn_R_50_FPN_Xconv1fc_1x_gn.yaml')
cfg.merge_from_file(args.config_file)
cfg.defrost()
cfg.OUTPUT_DIR = os.path.join(info['training_dir'], args.sub_dataset)
cfg.MODEL.WEIGHT = os.path.join(info['dataset_dir'], info['experiment'],
'Detector',
'Iter{}.pth'.format(info['iter']))
cfg.SOLVER.IMS_PER_BATCH = num_gpus * 4
cfg.TEST.IMS_PER_BATCH = num_gpus * 16
cfg.SOLVER.BASE_LR = 0.002
cfg.freeze()
mkdir(cfg.OUTPUT_DIR)
if args.sub_dataset is None:
args.sub_dataset = ""
if args.vis_title is None:
args.vis_title = os.path.basename(cfg.OUTPUT_DIR)
logger = setup_logger("maskrcnn_benchmark", cfg.OUTPUT_DIR, get_rank())
logger.info("Using {} GPUs".format(num_gpus))
logger.info(args)
logger.info("Collecting env info (might take some time)")
logger.info("\n" + collect_env_info())
DatasetCatalog = None
train_dataset = cfg.DATASETS.TRAIN[0]
test_dataset = cfg.DATASETS.TEST[0]
paths_catalog = import_file("maskrcnn_benchmark.config.paths_catalog",
cfg.PATHS_CATALOG, True)
if args.sub_dataset != "":
DatasetCatalog = paths_catalog.DatasetCatalog
DatasetCatalog.DATASETS[train_dataset]['img_dir'] = os.path.join(
info['dataset_dir'], 'Images')
DatasetCatalog.DATASETS[train_dataset]['ann_file'] = os.path.join(
info['dataset_dir'], 'RCNN_data', 'train.json')
DatasetCatalog.DATASETS[test_dataset]['img_dir'] = os.path.join(
info['dataset_dir'], 'Images')
DatasetCatalog.DATASETS[test_dataset]['ann_file'] = os.path.join(
info['dataset_dir'], 'RCNN_data', 'test.json')
data = json.load(
open(DatasetCatalog.DATASETS[train_dataset]['ann_file']))
else:
data = json.load(
open(paths_catalog.DatasetCatalog.DATASETS[train_dataset]
['ann_file']))
iters_per_epoch = len(data['images'])
iters_per_epoch = math.ceil(iters_per_epoch / cfg.SOLVER.IMS_PER_BATCH)
args.iters_per_epoch = iters_per_epoch
cfg.defrost()
cfg.SOLVER.MAX_ITER = round(args.epochs * args.scale * iters_per_epoch)
cfg.SOLVER.STEPS = (round(8 * args.scale * iters_per_epoch),
round(11 * args.scale * iters_per_epoch),
round(16 * args.scale * iters_per_epoch))
cfg.freeze()
logger.info("Loaded configuration file {}".format(args.config_file))
with open(args.config_file, "r") as cf:
config_str = "\n" + cf.read()
logger.info(config_str)
logger.info("Running with config:\n{}".format(cfg))
logger.info(DatasetCatalog)
output_config_path = os.path.join(cfg.OUTPUT_DIR, 'config.yml')
logger.info("Saving config into: {}".format(output_config_path))
# save overloaded model config in the output directory
save_config(cfg, output_config_path)
if args.train:
args.skip_train = False
model = network.train(cfg, args, DatasetCatalog)
if args.test:
network.test(cfg, args, model=None, DatasetCatalog=DatasetCatalog)
