示例#1
0
    def _construct_and_fill_model(self):
        # Progress reporting to show a progress bar in the UI.
        model_build_progress = sly.Progress('Building model:', 1)

        # Check the class name --> index mapping to infer the number of model output dimensions.
        num_classes = max(self.class_title_to_idx.values()) + 1

        # Initialize the model.
        model = PyTorchSegmentation(num_classes=num_classes)
        sly.logger.info('Model has been instantiated.')

        # Load model weights appropriate for the given training mode.
        weights_rw = WeightsRW(sly.TaskPaths.MODEL_DIR)
        weights_init_type = self.config[WEIGHTS_INIT_TYPE]
        if weights_init_type == TRANSFER_LEARNING:
            # For transfer learning, do not attempt to load the weights for the model head. The existing snapshot may
            # have been trained on a different dataset, even on a different set of classes, and is in general not
            # compatible with the current model even in terms of dimensions. The head of the model will be initialized
            # randomly.
            self._model = weights_rw.load_for_transfer_learning(
                model, ignore_matching_layers=['_head'], logger=sly.logger)
        elif weights_init_type == CONTINUE_TRAINING:
            # Continuing training from an older snapshot requires full compatibility between the two models, including
            # class index mapping. Hence the snapshot weights must exactly match the structure of our model instance.
            self._model = weights_rw.load_strictly(model)

        # Model weights have been loaded, move them over to the GPU.
        self._model.cuda()

        # Advance the progress bar and log a progress message.
        sly.logger.info('Weights have been loaded.',
                        extra={WEIGHTS_INIT_TYPE: weights_init_type})
        model_build_progress.iter_done_report()
示例#2
0
    def _construct_and_fill_model(self):
        # TODO: Move it progress to base class
        progress_dummy = sly.Progress('Building model:', 1)
        progress_dummy.iter_done_report()
        self.model = create_model(
            n_cls=(max(self.class_title_to_idx.values()) + 1),
            device_ids=self.device_ids)

        if sly.fs.dir_empty(sly.TaskPaths.MODEL_DIR):
            sly.logger.info('Weights will not be inited.')
            # @TODO: add random init (m.weight.data.normal_(0, math.sqrt(2. / n))
        else:
            wi_type = self.config['weights_init_type']
            ewit = {'weights_init_type': wi_type}
            sly.logger.info('Weights will be inited from given model.',
                            extra=ewit)

            weights_rw = WeightsRW(sly.TaskPaths.MODEL_DIR)
            if wi_type == TRANSFER_LEARNING:
                self.model = weights_rw.load_for_transfer_learning(
                    self.model,
                    ignore_matching_layers=['last_conv'],
                    logger=logger)
            elif wi_type == CONTINUE_TRAINING:
                self.model = weights_rw.load_strictly(self.model)

            sly.logger.info('Weights are loaded.', extra=ewit)
示例#3
0
    def _construct_and_fill_model(self):
        super()._construct_and_fill_model()
        device_ids = sly.env.remap_gpu_devices([self._config[GPU_DEVICE]])
        num_layers = determine_resnet_model_configuration(TaskPaths.MODEL_CONFIG_PATH)
        self.model = create_model(num_layers=num_layers, n_cls=(max(self.classification_tags_to_idx.values()) + 1),
                                  device_ids=device_ids)

        self.model = WeightsRW(TaskPaths.MODEL_DIR).load_strictly(self.model)
        self.model.eval()
        logger.info('Weights are loaded.')
示例#4
0
    def _construct_and_fill_model(self):
        progress_dummy = sly.Progress('Building model:', 1)
        progress_dummy.iter_done_report()

        self.model = create_model(self.num_layers,
                                  n_cls=len(self.classification_tags_sorted), device_ids=self.device_ids)

        if sly.fs.dir_empty(sly.TaskPaths.MODEL_DIR):
            logger.info('Weights will not be inited.')
            # @TODO: add random init (m.weight.data.normal_(0, math.sqrt(2. / n))
        else:
            wi_type = self.config['weights_init_type']
            ewit = {'weights_init_type': wi_type}
            logger.info('Weights will be inited from given model.', extra=ewit)

            weights_rw = WeightsRW(sly.TaskPaths.MODEL_DIR)
            if wi_type == TRANSFER_LEARNING:
                self.model = weights_rw.load_for_transfer_learning(self.model, ignore_matching_layers=['fc'],
                                                                   logger=logger)
            elif wi_type == CONTINUE_TRAINING:
                self.model = weights_rw.load_strictly(self.model)

            logger.info('Weights are loaded.', extra=ewit)
示例#5
0
    def _construct_and_fill_model(self):
        super()._construct_and_fill_model()

        self.device_ids = sly.env.remap_gpu_devices([self._config[GPU_DEVICE]])
        n_cls = (max(self.out_class_mapping.keys()) + 1)
        use_batchnorm = self.train_config[SETTINGS]['use_batchnorm']

        model = ICNet(n_classes=n_cls,
                      input_size=self.input_size,
                      is_batchnorm=use_batchnorm)
        WeightsRW(TaskPaths.MODEL_DIR).load_strictly(model)
        logger.info('Weights are loaded.')
        self.model = DataParallel(model, device_ids=self.device_ids)
        self.model.eval()
        self.model.cuda()
示例#6
0
    def _construct_and_fill_model(self):
        super()._construct_and_fill_model()   # Progress reporting done by the base class.
        # Check the class index --> name mapping to infer the number of model output dimensions.
        num_classes = max(self.out_class_mapping.keys()) + 1

        # Initialize the model.
        self._model = PyTorchSegmentation(num_classes=num_classes)
        sly.logger.info('Model has been instantiated.')

        # Load model weights.
        WeightsRW(sly.TaskPaths.MODEL_DIR).load_strictly(self._model)

        # Switch the model into evaluation mode (disable gradients computation and batchnorm updates).
        self._model.eval()

        # Move the model to the GPU.
        self._model.cuda()
        sly.logger.info('Model weights have been loaded.')
示例#7
0
 def _dump_model_weights(self, out_dir):
     WeightsRW(out_dir).save(self.model)
示例#8
0
 def _dump_model_weights(self, out_dir):
     # Framework-specific logic to snapshot the model weights.
     WeightsRW(out_dir).save(self._model)
示例#9
0
class ResnetSingleImageApplier(SingleImageInferenceBase):
    @property
    def classification_tags_key(self):
        return config_lib.classification_tags_key()

    @property
    def classification_tags_to_idx_key(self):
        return config_lib.classification_tags_to_idx_key()

    @property
    def train_classes_key(self):
        return config_lib.train_classes_key()

    @property
    def class_title_to_idx_key(self):
        return config_lib.class_to_idx_config_key()

    def _model_out_tags(self):
        temp_collection = TagMetaCollection.from_json(
            self.train_config[self.classification_tags_key])
        res_collection = TagMetaCollection([
            TagMeta(x.name, TagValueType.ANY_NUMBER) for x in temp_collection
        ])
        return res_collection

    def _load_train_config(
            self):  # @TODO: partly copypasted from SingleImageInferenceBase
        self._load_raw_model_config_json()

        self.classification_tags = self._model_out_tags()
        logger.info('Read model out tags',
                    extra={'tags': self.classification_tags.to_json()})
        self.classification_tags_to_idx = self.train_config[
            self.classification_tags_to_idx_key]
        logger.info('Read model internal tags mapping',
                    extra={'tags_mapping': self.classification_tags_to_idx})

        self._model_out_meta = ProjectMeta(obj_classes=ObjClassCollection(),
                                           tag_metas=self.classification_tags)

        self.idx_to_classification_tags = {
            v: k
            for k, v in self.classification_tags_to_idx.items()
        }
        self._determine_model_input_size()

    def _validate_model_config(self, config):
        JsonConfigValidator().validate_inference_cfg(config)

    @staticmethod
    def get_default_config():
        return {GPU_DEVICE: 0}

    def _construct_and_fill_model(self):
        super()._construct_and_fill_model()
        device_ids = sly.env.remap_gpu_devices([self._config[GPU_DEVICE]])
        num_layers = determine_resnet_model_configuration(
            TaskPaths.MODEL_CONFIG_PATH)
        self.model = create_model(
            num_layers=num_layers,
            n_cls=(max(self.classification_tags_to_idx.values()) + 1),
            device_ids=device_ids)

        self.model = WeightsRW(TaskPaths.MODEL_DIR).load_strictly(self.model)
        self.model.eval()
        logger.info('Weights are loaded.')

    def inference(self, img, ann):
        output = infer_on_img(img, self.input_size, self.model)
        tag_id = np.argmax(output)
        score = output[tag_id]
        tag_name = self.idx_to_classification_tags[tag_id]
        tag = Tag(self.classification_tags.get(tag_name),
                  round(float(score), 4))
        tags = TagCollection([tag])
        return Annotation(ann.img_size, img_tags=tags)
示例#10
0
def create_model_for_inference(n_cls, device_ids, model_dir):
    model = create_model(n_cls, device_ids)
    model = WeightsRW(model_dir).load_strictly(model)
    model.eval()
    return model