| 22 | |
| 23 | |
| 24 | class Inference(ClsSolver): |
| 25 | def __init__(self, config): |
| 26 | self.image_dir = config["image_dir"] |
| 27 | self.meta_file = config.get("meta_file", "") |
| 28 | |
| 29 | self.output = config.get("output", "inference_results") |
| 30 | self.recover = config.get("recover", "") |
| 31 | self.cam = config.get("cam", False) |
| 32 | self.visualize = config.get("visualize", False) |
| 33 | self.sample = config.get("sample", -1) |
| 34 | self.feature_name = config.get("name", "module.layer4") |
| 35 | if "module" not in self.feature_name: |
| 36 | self.feature_name = "module." + self.feature_name |
| 37 | self.feature = None |
| 38 | self.gradient = None |
| 39 | super(Inference, self).__init__(config["config"]) |
| 40 | |
| 41 | def setup_env(self): |
| 42 | # dist |
| 43 | self.dist = EasyDict() |
| 44 | self.dist.rank, self.dist.world_size = link.get_rank(), link.get_world_size() |
| 45 | self.prototype_info.world_size = self.dist.world_size |
| 46 | # directories |
| 47 | self.path = EasyDict() |
| 48 | self.path.root_path = os.path.dirname(self.config_file) |
| 49 | self.path.result_path = os.path.abspath(self.output) |
| 50 | makedir(self.path.result_path) |
| 51 | # logger |
| 52 | create_logger(os.path.join(self.path.root_path, 'log.txt')) |
| 53 | self.logger = get_logger(__name__) |
| 54 | self.logger.info(f'config: {pprint.pformat(self.config)}') |
| 55 | if 'SLURM_NODELIST' in os.environ: |
| 56 | self.logger.info(f"hostnames: {os.environ['SLURM_NODELIST']}") |
| 57 | # load pretrain checkpoint |
| 58 | if self.recover != "": |
| 59 | self.state = torch.load(self.recover, 'cpu') |
| 60 | self.logger.info(f"Recovering from {self.recover}, keys={list(self.state.keys())}") |
| 61 | |
| 62 | elif hasattr(self.config.saver, 'pretrain'): |
| 63 | self.state = torch.load(self.config.saver.pretrain.path, 'cpu') |
| 64 | self.logger.info(f"Recovering from {self.config.saver.pretrain.path}, keys={list(self.state.keys())}") |
| 65 | if hasattr(self.config.saver.pretrain, 'ignore'): |
| 66 | self.state = modify_state(self.state, self.config.saver.pretrain.ignore) |
| 67 | else: |
| 68 | self.state = {} |
| 69 | self.state['last_iter'] = 0 |
| 70 | # others |
| 71 | torch.backends.cudnn.benchmark = True |
| 72 | |
| 73 | def build_data(self): |
| 74 | self.config.data.max_iter = self.config.lr_scheduler.kwargs.max_iter |
| 75 | self.config.data.last_iter = self.state['last_iter'] |
| 76 | |
| 77 | root_dir, input_file = self.generate_custom_data() |
| 78 | self.config.data.test.root_dir = root_dir |
| 79 | self.config.data.test.meta_file = input_file |
| 80 | |
| 81 | if self.config.data.get('type', 'imagenet') == 'imagenet': |