-
Notifications
You must be signed in to change notification settings - Fork 0
/
feature_extraction.py
54 lines (45 loc) · 2.28 KB
/
feature_extraction.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
__author__ = "shekkizh"
"""Modified code for feature extraction using VISSL tutorial and tools"""
import argparse
import os
from typing import Any, List
from vissl.utils.hydra_config import convert_to_attrdict, is_hydra_available
from hydra.experimental import compose, initialize_config_module
from vissl.utils.distributed_launcher import launch_distributed
from vissl.hooks import default_hook_generator
# from vissl.utils.env import set_env_vars
from vissl.data.dataset_catalog import VisslDatasetCatalog
parser = argparse.ArgumentParser(description='VISSL extract features')
parser.add_argument('--model_url',
default='https://dl.fbaipublicfiles.com/vissl/model_zoo/deepclusterv2_800ep_pretrain.pth.tar',
help='Model to download - https://github.com/facebookresearch/vissl/blob/master/MODEL_ZOO.md')
parser.add_argument('--logs_dir', default='/scratch/shekkizh/logs/VISSL')
parser.add_argument("--config", default="imagenet1k_resnet50_trunk_features.yaml",
help="config file to extract features")
def hydra_main(overrides: List[Any]):
print(f"####### overrides: {overrides}")
with initialize_config_module(config_module="vissl.config"):
cfg = compose("defaults", overrides=overrides)
args, config = convert_to_attrdict(cfg)
# set_env_vars(local_rank=0, node_id=0, cfg=config)
launch_distributed(
cfg=config,
node_id=args.node_id,
engine_name=args.engine_name,
hook_generator=default_hook_generator,
)
if __name__ == "__main__":
args = parser.parse_args()
print("Retrieving model weights from VISSL MODEL ZOO")
basename = os.path.basename(args.model_url)
weights_file = os.path.join('/scratch/shekkizh/torch_hub/checkpoints/', basename)
if not os.path.exists(weights_file):
os.system(f"wget -O {weights_file} -L {args.model_url}")
logs_dir = os.path.join(args.logs_dir, basename.split('.')[0])
# print imagenet path
print(VisslDatasetCatalog.get("imagenet1k_folder"))
overrides = [f"config={args.config}", f"config.CHECKPOINT.DIR={logs_dir}",
f"config.MODEL.WEIGHTS_INIT.PARAMS_FILE={weights_file}"]
assert is_hydra_available(), "Make sure to install hydra"
overrides.append("hydra.verbose=true")
hydra_main(overrides=overrides)