def __init__(self,
model,
save_dir="",
*,
save_to_disk=None,
**checkpointables):
is_main_process = comm.is_main_process()
super().__init__(
model,
save_dir,
save_to_disk=is_main_process
if save_to_disk is None else save_to_disk,
**checkpointables,
)
if hasattr(self, "path_manager"):
self.path_manager = PathManager
else:
# This could only happen for open source
# TODO remove after upgrading fvcore
from fvcore.common.file_io import PathManager as g_PathManager
for handler in PathManager._path_handlers.values():
try:
g_PathManager.register_handler(handler)
except KeyError:
pass
def test_failure(self) -> None:
self.assertFalse(PathManager.exists("test://foo"))
PathManager.register_handler(TestPathHandler())
# make sure that TestPathHandler is being used
self.assertTrue(PathManager.exists("test://foo"))
checkpoint_folder = "test://root"
checkpoint_hook = CheckpointHook(checkpoint_folder, {}, phase_types=["train"])
config = get_test_task_config()
task = build_task(config)
task.prepare()
# we should raise an exception while trying to save the checkpoint
with self.assertRaises(TestException):
checkpoint_hook.on_phase_end(task)
def set_env_vars(local_rank: int, node_id: int, cfg: AttrDict):
"""
Set some environment variables like total number of gpus used in training,
distributed rank and local rank of the current gpu, whether to print the
nccl debugging info and tuning nccl settings.
"""
os.environ["WORLD_SIZE"] = str(
cfg.DISTRIBUTED.NUM_NODES * cfg.DISTRIBUTED.NUM_PROC_PER_NODE
)
dist_rank = cfg.DISTRIBUTED.NUM_PROC_PER_NODE * node_id + local_rank
os.environ["RANK"] = str(dist_rank)
os.environ["LOCAL_RANK"] = str(local_rank)
if cfg.DISTRIBUTED.NCCL_DEBUG:
os.environ["NCCL_DEBUG"] = "INFO"
if cfg.DISTRIBUTED.NCCL_SOCKET_NTHREADS:
logging.info(
f"local_rank: {local_rank}, "
f"using NCCL_SOCKET_NTHREADS: {cfg.DISTRIBUTED.NCCL_SOCKET_NTHREADS}"
)
os.environ["NCCL_SOCKET_NTHREADS"] = str(cfg.DISTRIBUTED.NCCL_SOCKET_NTHREADS)
# register http handler to support reading the urls
PathManager.register_handler(HTTPURLHandler(), allow_override=True)
def register_handler(handler) -> None:
if FVCorePathManager:
return FVCorePathManager.register_handler(handler=handler)
return [self.PREFIX]
def _get_local_path(self, path):
logger = logging.getLogger(__name__)
catalog_path = ModelZooURL.get(self.PREFIX, path[len(self.PREFIX):])
logger.info("Catalog entry {} points to {}".format(path, catalog_path))
return PathManager.get_local_path(catalog_path)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(self._get_local_path(path), mode, **kwargs)
class LocalModelCatalogHandler(PathHandler):
PREFIX = "local://"
def _get_supported_prefixes(self):
return [self.PREFIX]
def _get_local_path(self, path):
logger = logging.getLogger(__name__)
catalog_path = ModelZooURL.get(self.PREFIX, path[len(self.PREFIX):])
logger.info("Catalog entry {} points to {}".format(path, catalog_path))
return PathManager.get_local_path(catalog_path)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(self._get_local_path(path), mode, **kwargs)
PathManager.register_handler(TorchvisionModelCatalogHandler())
PathManager.register_handler(LocalModelCatalogHandler())
import os
import shutil
from typing import List, Optional
logger = logging.getLogger(__file__)
try:
from fvcore.common.file_io import PathManager as FVCorePathManager
try:
# [FB only - for now] AWS PathHandler for PathManager
from .fb_pathhandlers import S3PathHandler
FVCorePathManager.register_handler(S3PathHandler())
except KeyError:
logging.warning("S3PathHandler already registered.")
except ImportError:
logging.debug(
"S3PathHandler couldn't be imported. Either missing fb-only files, or boto3 module."
)
except ImportError:
FVCorePathManager = None
class PathManager:
"""
Wrapper for insulating OSS I/O (using Python builtin operations) from
fvcore's PathManager abstraction (for transparently handling various
# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
from fvcore.common.file_io import PathHandler, PathManager
__all__ = ["ShapenetHandler"]
class ShapenetHandler(PathHandler):
"""
Resolve anything that's in Mesh R-CNN's model zoo.
"""
PREFIX = "shapenet://"
SHAPENET_PREFIX = "https://dl.fbaipublicfiles.com/meshrcnn/shapenet/"
def _get_supported_prefixes(self):
return [self.PREFIX]
def _get_local_path(self, path):
name = path[len(self.PREFIX):]
return PathManager.get_local_path(self.SHAPENET_PREFIX + name)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(self._get_local_path(path), mode, **kwargs)
PathManager.register_handler(ShapenetHandler())
# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
from fvcore.common.file_io import PathHandler, PathManager
__all__ = ["MeshRCNNHandler"]
class MeshRCNNHandler(PathHandler):
"""
Resolve anything that's in Mesh R-CNN's model zoo.
"""
PREFIX = "meshrcnn://"
MESHRCNN_PREFIX = "https://dl.fbaipublicfiles.com/meshrcnn/pix3d/"
def _get_supported_prefixes(self):
return [self.PREFIX]
def _get_local_path(self, path):
name = path[len(self.PREFIX):]
return PathManager.get_local_path(self.MESHRCNN_PREFIX + name)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(self._get_local_path(path), mode, **kwargs)
PathManager.register_handler(MeshRCNNHandler())
logger = logging.getLogger(__name__)
catalog_path = ModelCatalog.get(path[len(self.PREFIX) :])
logger.info("Catalog entry {} points to {}".format(path, catalog_path))
return PathManager.get_local_path(catalog_path)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(self._get_local_path(path), mode, **kwargs)
class Detectron2Handler(PathHandler):
"""
Resolve anything that's in Detectron2 model zoo.
"""
PREFIX = "detectron2://"
S3_DETECTRON2_PREFIX = "https://dl.fbaipublicfiles.com/detectron2/"
def _get_supported_prefixes(self):
return [self.PREFIX]
def _get_local_path(self, path):
name = path[len(self.PREFIX) :]
return PathManager.get_local_path(self.S3_DETECTRON2_PREFIX + name)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(self._get_local_path(path), mode, **kwargs)
PathManager.register_handler(ModelCatalogHandler())
PathManager.register_handler(Detectron2Handler())
class LayoutParserHandler(PathHandler):
"""
Resolve anything that's in LayoutParser model zoo.
"""
PREFIX = "lp://"
def _get_supported_prefixes(self):
return [self.PREFIX]
def _get_local_path(self, path):
model_name = path[len(self.PREFIX):]
dataset_name, *model_name, data_type = model_name.split('/')
if data_type == 'weight':
model_url = MODEL_CATALOG[dataset_name]['/'.join(model_name)]
elif data_type == 'config':
model_url = CONFIG_CATALOG[dataset_name]['/'.join(model_name)]
else:
raise ValueError(f"Unknown data_type {data_type}")
return PathManager.get_local_path(model_url)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(self._get_local_path(path), mode, **kwargs)
PathManager.register_handler(DropboxHandler())
PathManager.register_handler(LayoutParserHandler())
def register_http_url_handler():
"""
support reading file from url starting with "http://", "https://", "ftp://"
"""
PathManager.register_handler(HTTPURLHandler(), allow_override=True)
im = cv2.imdecode(np.frombuffer(data, np.uint8), flag)
return im
@staticmethod
def read(path):
zip_path, path_img = ZipReader.split_zip_style_path(path)
suffix = zip_path.split('/')[-1][:-4]
path_img = suffix + '/' + path_img
zfile = ZipReader.get_zipfile(zip_path)
data = zfile.read(path_img)
return data
class ZipHandler(PathHandler):
PREFIX = "zip://"
def _get_supported_prefixes(self):
return ["zip://"]
def _get_local_path(self, path):
name = path[len(self.PREFIX) :]
return name
@contextmanager
def _open(self, path, mode="r", **kwargs):
yield self._get_local_path(path)
PathManager.register_handler(ZipHandler())
"""FS-Det Model Handler."""
from fvcore.common.file_io import PathHandler, PathManager
import logging
class FsDetHandler(PathHandler):
"""
Resolve anything that's in FsDet model zoo.
"""
PREFIX = "fsdet://"
URL_PREFIX = "http://dl.yf.io/fs-det/models/"
def _get_supported_prefixes(self):
return [self.PREFIX]
def _get_local_path(self, path):
name = path[len(self.PREFIX):]
return PathManager.get_local_path(self.URL_PREFIX + name)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(self._get_local_path(path), mode, **kwargs)
PathManager.register_handler(FsDetHandler())
logger = logging.getLogger(__name__)
catalog_path = ModelCatalog.get(path[len(self.PREFIX) :])
logger.info("Catalog entry {} points to {}".format(path, catalog_path))
return PathManager.get_local_path(catalog_path)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(self._get_local_path(path), mode, **kwargs)
class mydlHandler(PathHandler):
"""
Resolve anything that's in mydl model zoo.
"""
PREFIX = "mydl://"
S3_DETECTRON2_PREFIX = "https://dl.fbaipublicfiles.com/mydl/"
def _get_supported_prefixes(self):
return [self.PREFIX]
def _get_local_path(self, path):
name = path[len(self.PREFIX) :]
return PathManager.get_local_path(self.S3_DETECTRON2_PREFIX + name)
def _open(self, path, mode="r", **kwargs):
return PathManager.open(self._get_local_path(path), mode, **kwargs)
PathManager.register_handler(ModelCatalogHandler())
PathManager.register_handler(mydlHandler())