def create_and_check_for_image_classification(self, config, pixel_values, labels): config.num_labels = self.type_sequence_label_size model = TFViTForImageClassification(config) result = model(pixel_values, labels=labels, training=False) self.parent.assertEqual(result.logits.shape, (self.batch_size, self.type_sequence_label_size)) # Test with an image with different size than the one specified in config. image_size = self.image_size // 2 pixel_values = pixel_values[:, :, :image_size, :image_size] result = model(pixel_values, interpolate_pos_encoding=True, training=False) self.parent.assertEqual(result.logits.shape, (self.batch_size, self.type_sequence_label_size)) # test greyscale images config.num_channels = 1 model = TFViTForImageClassification(config) pixel_values = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]) result = model(pixel_values) self.parent.assertEqual(result.logits.shape, (self.batch_size, self.type_sequence_label_size))
def test_inference_image_classification_head(self): model = TFViTForImageClassification.from_pretrained("google/vit-base-patch16-224") feature_extractor = self.default_feature_extractor image = prepare_img() inputs = feature_extractor(images=image, return_tensors="tf") # forward pass outputs = model(**inputs) # verify the logits expected_shape = tf.TensorShape((1, 1000)) self.assertEqual(outputs.logits.shape, expected_shape) expected_slice = tf.constant([-0.2744, 0.8215, -0.0836]) tf.debugging.assert_near(outputs.logits[0, :3], expected_slice, atol=1e-4)
dataset_train = dataset_train.map(load_and_process_image, num_parallel_calls=tf.data.AUTOTUNE) dataset_train = dataset_train.shuffle(len(dataset_train)).batch( BATCH_SIZE, drop_remainder=True) dataset_validation = tf.data.Dataset.from_tensor_slices( (image_paths['validation'], image_labels['validation'])) dataset_validation = dataset_validation.map( load_and_process_image, num_parallel_calls=tf.data.AUTOTUNE) dataset_validation = dataset_validation.batch(BATCH_SIZE, drop_remainder=True) # ## Model # # ### Initialization model = TFViTForImageClassification.from_pretrained( VITMODEL, num_labels=1, ignore_mismatched_sizes=True) LR = 1e-5 optimizer = tf.keras.optimizers.Adam(learning_rate=LR) loss = tf.keras.losses.BinaryCrossentropy(from_logits=False) metric = 'accuracy' model.compile(optimizer=optimizer, loss=loss, metrics=[metric]) print(model.summary()) # ### Learning logdir = os.path.join( os.getcwd(), "logs",