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()
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)
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 _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)
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()
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.')
def _dump_model_weights(self, out_dir):
WeightsRW(out_dir).save(self.model)
def _dump_model_weights(self, out_dir):
# Framework-specific logic to snapshot the model weights.
WeightsRW(out_dir).save(self._model)
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)
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
