def test_model_creation(self, project_dim, num_proj_layers, ft_proj_idx): input_size = 224 inputs = np.random.rand(2, input_size, input_size, 3) input_specs = tf.keras.layers.InputSpec( shape=[None, input_size, input_size, 3]) tf.keras.backend.set_image_data_format('channels_last') backbone = backbones.ResNet(model_id=50, activation='relu', input_specs=input_specs) projection_head = simclr_head.ProjectionHead( proj_output_dim=project_dim, num_proj_layers=num_proj_layers, ft_proj_idx=ft_proj_idx ) num_classes = 10 supervised_head = simclr_head.ClassificationHead( num_classes=10 ) model = simclr_model.SimCLRModel( input_specs=input_specs, backbone=backbone, projection_head=projection_head, supervised_head=supervised_head, mode=simclr_model.PRETRAIN ) outputs = model(inputs) projection_outputs = outputs[simclr_model.PROJECTION_OUTPUT_KEY] supervised_outputs = outputs[simclr_model.SUPERVISED_OUTPUT_KEY] self.assertAllEqual(projection_outputs.shape.as_list(), [2, project_dim]) self.assertAllEqual([2, num_classes], supervised_outputs.numpy().shape)
def _instantiate_sub_tasks(self) -> Dict[Text, tf.keras.Model]: tasks = {} for model_config in self._config.heads: # Build supervised head supervised_head_config = model_config.supervised_head if supervised_head_config: if supervised_head_config.zero_init: s_kernel_initializer = 'zeros' else: s_kernel_initializer = 'random_uniform' supervised_head = simclr_head.ClassificationHead( num_classes=supervised_head_config.num_classes, kernel_initializer=s_kernel_initializer, kernel_regularizer=self._l2_regularizer) else: supervised_head = None tasks[model_config.task_name] = simclr_model.SimCLRModel( input_specs=self._input_specs, backbone=self._backbone, projection_head=self._projection_head, supervised_head=supervised_head, mode=model_config.mode, backbone_trainable=self._config.backbone_trainable) return tasks
def test_head_creation(self, num_classes): test_layer = simclr_head.ClassificationHead(num_classes=num_classes) input_dim = 64 x = tf.keras.Input(shape=(input_dim, )) out_x = test_layer(x) self.assertAllEqual(out_x.shape.as_list(), [None, num_classes])
def build_model(self): model_config = self.task_config.model input_specs = tf.keras.layers.InputSpec(shape=[None] + model_config.input_size) l2_weight_decay = self.task_config.loss.l2_weight_decay # Divide weight decay by 2.0 to match the implementation of tf.nn.l2_loss. # (https://www.tensorflow.org/api_docs/python/tf/keras/regularizers/l2) # (https://www.tensorflow.org/api_docs/python/tf/nn/l2_loss) l2_regularizer = ( tf.keras.regularizers.l2(l2_weight_decay / 2.0) if l2_weight_decay else None) # Build backbone backbone = backbones.factory.build_backbone( input_specs=input_specs, backbone_config=model_config.backbone, norm_activation_config=model_config.norm_activation, l2_regularizer=l2_regularizer) # Build projection head norm_activation_config = model_config.norm_activation projection_head_config = model_config.projection_head projection_head = simclr_head.ProjectionHead( proj_output_dim=projection_head_config.proj_output_dim, num_proj_layers=projection_head_config.num_proj_layers, ft_proj_idx=projection_head_config.ft_proj_idx, kernel_regularizer=l2_regularizer, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon) # Build supervised head supervised_head_config = model_config.supervised_head if supervised_head_config: if supervised_head_config.zero_init: s_kernel_initializer = 'zeros' else: s_kernel_initializer = 'random_uniform' supervised_head = simclr_head.ClassificationHead( num_classes=supervised_head_config.num_classes, kernel_initializer=s_kernel_initializer, kernel_regularizer=l2_regularizer) else: supervised_head = None model = simclr_model.SimCLRModel( input_specs=input_specs, backbone=backbone, projection_head=projection_head, supervised_head=supervised_head, mode=model_config.mode, backbone_trainable=model_config.backbone_trainable) logging.info(model.get_config()) return model
def _instantiate_sub_tasks(self) -> Dict[Text, tf.keras.Model]: tasks = {} # Build the shared projection head norm_activation_config = self._config.norm_activation projection_head_config = self._config.projection_head projection_head = simclr_head.ProjectionHead( proj_output_dim=projection_head_config.proj_output_dim, num_proj_layers=projection_head_config.num_proj_layers, ft_proj_idx=projection_head_config.ft_proj_idx, kernel_regularizer=self._l2_regularizer, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon) for model_config in self._config.heads: # Build supervised head supervised_head_config = model_config.supervised_head if supervised_head_config: if supervised_head_config.zero_init: s_kernel_initializer = 'zeros' else: s_kernel_initializer = 'random_uniform' supervised_head = simclr_head.ClassificationHead( num_classes=supervised_head_config.num_classes, kernel_initializer=s_kernel_initializer, kernel_regularizer=self._l2_regularizer) else: supervised_head = None tasks[model_config.mode] = simclr_model.SimCLRModel( input_specs=self._input_specs, backbone=self._backbone, projection_head=projection_head, supervised_head=supervised_head, mode=model_config.mode, backbone_trainable=self._config.backbone_trainable) return tasks