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",