metadata

license: apache-2.0
tags:
  - generated_from_trainer
datasets:
  - chest-xray-classification
  - keremberke/chest-xray-classification
metrics:
  - accuracy
model-index:
  - name: vit-xray-pneumonia-classification
    results:
      - task:
          name: Image Classification
          type: image-classification
        dataset:
          name: chest-xray-classification
          type: chest-xray-classification
          config: full
          split: validation
          args: full
        metrics:
          - name: Accuracy
            type: accuracy
            value: 0.9742489270386266
pipeline_tag: image-classification

vit-xray-pneumonia-classification

This model is a fine-tuned version of google/vit-base-patch16-224-in21k on the chest-xray-classification dataset. It achieves the following results on the evaluation set:

Loss: 0.0868
Accuracy: 0.9742

Inference example

from transformers import pipeline

classifier = pipeline(model="lxyuan/vit-xray-pneumonia-classification")

# image taken from https://www.news-medical.net/health/What-is-Viral-Pneumonia.aspx
classifier("https://d2jx2rerrg6sh3.cloudfront.net/image-handler/ts/20200618040600/ri/650/picture/2020/6/shutterstock_786937069.jpg")

>>>
[{'score': 0.990334689617157, 'label': 'PNEUMONIA'},
 {'score': 0.009665317833423615, 'label': 'NORMAL'}]

Training procedure

Notebook link: here

Training hyperparameters

The following hyperparameters were used during training:

learning_rate: 5e-05
train_batch_size: 16
eval_batch_size: 16
seed: 42
gradient_accumulation_steps: 4
total_train_batch_size: 64
optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
lr_scheduler_type: linear
lr_scheduler_warmup_ratio: 0.1
num_epochs: 15

from transformers import EarlyStoppingCallback

training_args = TrainingArguments(
    output_dir="vit-xray-pneumonia-classification",
    remove_unused_columns=False,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    logging_strategy="epoch",
    learning_rate=5e-5,
    per_device_train_batch_size=16,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=16,
    num_train_epochs=15,
    save_total_limit=2,
    warmup_ratio=0.1,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    fp16=True,
    push_to_hub=True,
    report_to="tensorboard"
)

early_stopping = EarlyStoppingCallback(early_stopping_patience=3)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    tokenizer=processor,
    compute_metrics=compute_metrics,
    callbacks=[early_stopping],
)

Training results

Training Loss	Epoch	Step	Validation Loss	Accuracy
0.5152	0.99	63	0.2507	0.9245
0.2334	1.99	127	0.1766	0.9382
0.1647	3.0	191	0.1218	0.9588
0.144	4.0	255	0.1222	0.9502
0.1348	4.99	318	0.1293	0.9571
0.1276	5.99	382	0.1000	0.9665
0.1175	7.0	446	0.1177	0.9502
0.109	8.0	510	0.1079	0.9665
0.0914	8.99	573	0.0804	0.9717
0.0872	9.99	637	0.0800	0.9717
0.0804	11.0	701	0.0862	0.9682
0.0935	12.0	765	0.0883	0.9657
0.0686	12.99	828	0.0868	0.9742

Framework versions

Transformers 4.30.2
Pytorch 1.9.0+cu102
Datasets 2.12.0
Tokenizers 0.13.3