def _make_yolo_train_config(self):
src_size = self.config['input_size']
input_size = (src_size['height'], src_size['width'])
yolo_config = read_config(os.path.join(sly.TaskPaths.MODEL_DIR, MODEL_CFG))
[net_config] = [section for section in yolo_config if section.name == NET_SECTION]
net_overrides = {
'height': input_size[0],
'width': input_size[1],
'batch': self.config['batch_size']['train'],
'subdivisions': self.config['subdivisions']['train'],
'learning_rate': self.config['lr']
}
replace_config_section_values(net_config, net_overrides)
for section_idx, section in enumerate(yolo_config):
if section.name == YOLO_SECTION:
no_preceding_conv_section_msg = ('Unexpectedly found {!r} section in the config without immediately '
'preceding {!r} section.'.format(YOLO_SECTION, CONVOLUTIONAL_SECTION))
if section_idx == 0:
raise ValueError(no_preceding_conv_section_msg)
last_convolution = yolo_config[section_idx - 1]
if last_convolution.name != CONVOLUTIONAL_SECTION:
raise ValueError(no_preceding_conv_section_msg)
classes_item = find_data_item(section, 'classes')
classes_item[1] = len(self.out_classes)
filters_item = find_data_item(last_convolution, 'filters')
filters_item[1] = (len(self.out_classes) + 5) * 3
first_yolo_section_idx = next(idx for idx, section in enumerate(yolo_config) if section.name == YOLO_SECTION)
# Compute the index of the last layer to be loaded for transfer learning.
# This will be the last layer before any layers that depend on the number
# of predicted classes start.
# The layers that depend on the number of classes are YOLO layers, and their
# immediately preceding convolutional layers (since those have number of filters
# depending on the number of classes).
# So we find the first YOLO layer, subtract 2 to skip YOLO and preceding convolutional layer,
# and subtract another 2 to account for the first [net] and root sections of the config.
self._last_transfer_learning_load_layer_idx = first_yolo_section_idx - 4
self._effective_model_cfg_path = os.path.join('/tmp', MODEL_CFG)
write_config(yolo_config, self._effective_model_cfg_path)
logger.info('Model config created.')
def _construct_and_fill_model(self):
super()._construct_and_fill_model()
self.device_ids = sly.env.remap_gpu_devices([self._config[GPU_DEVICE]])
yolo_config = read_config(
os.path.join(sly.TaskPaths.MODEL_DIR, MODEL_CFG))
[net_config] = [
section for section in yolo_config if section.name == NET_SECTION
]
net_overrides = {'batch': 1, 'subdivisions': 1}
replace_config_section_values(net_config, net_overrides)
effective_model_cfg_path = os.path.join('/tmp', MODEL_CFG)
write_config(yolo_config, effective_model_cfg_path)
self.model = load_net(
effective_model_cfg_path.encode('utf-8'),
os.path.join(sly.TaskPaths.MODEL_DIR,
'model.weights').encode('utf-8'), 0)
sly.logger.info('Weights are loaded.')