|
import torch.nn as nn |
|
import numpy as np |
|
|
|
from transformers import ViTFeatureExtractor, ViTModel, ViTForImageClassification, TrainingArguments, Trainer, \ |
|
default_data_collator, EarlyStoppingCallback |
|
from transformers.modeling_outputs import SequenceClassifierOutput |
|
from datasets import load_dataset, load_metric, Features, ClassLabel, Array3D |
|
|
|
train_ds, test_ds = load_dataset('cifar10', split=['train[:5000]', 'test[:2000]']) |
|
splits = train_ds.train_test_split(test_size=0.1) |
|
train_ds = splits['train'] |
|
val_ds = splits['test'] |
|
|
|
feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k') |
|
data_collator = default_data_collator |
|
|
|
|
|
def preprocess_images(examples): |
|
images = examples['img'] |
|
images = [np.array(image, dtype=np.uint8) for image in images] |
|
images = [np.moveaxis(image, source=-1, destination=0) for image in images] |
|
inputs = feature_extractor(images=images) |
|
examples['pixel_values'] = inputs['pixel_values'] |
|
|
|
return examples |
|
|
|
|
|
features = Features({ |
|
'label': ClassLabel( |
|
names=['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']), |
|
'img': Array3D(dtype="int64", shape=(3, 32, 32)), |
|
'pixel_values': Array3D(dtype="float32", shape=(3, 224, 224)), |
|
}) |
|
|
|
preprocessed_train_ds = train_ds.map(preprocess_images, batched=True, features=features) |
|
preprocessed_val_ds = val_ds.map(preprocess_images, batched=True, features=features) |
|
preprocessed_test_ds = test_ds.map(preprocess_images, batched=True, features=features) |
|
|
|
|
|
class ViTForImageClassification2(nn.Module): |
|
def __init__(self, num_labels=10): |
|
super(ViTForImageClassification2, self).__init__() |
|
self.vit = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k') |
|
self.classifier = nn.Linear(self.vit.config.hidden_size, num_labels) |
|
self.num_labels = num_labels |
|
|
|
def forward(self, pixel_values, labels): |
|
outputs = self.vit(pixel_values=pixel_values) |
|
logits = self.classifier(outputs.last_hidden_state[:, 0]) |
|
|
|
loss = None |
|
if labels is not None: |
|
loss_fct = nn.CrossEntropyLoss() |
|
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) |
|
|
|
return SequenceClassifierOutput( |
|
loss=loss, |
|
logits=logits, |
|
hidden_states=outputs.hidden_states, |
|
attentions=outputs.attentions, |
|
) |
|
|
|
|
|
args = TrainingArguments( |
|
f"test-cifar-10", |
|
evaluation_strategy="epoch", |
|
learning_rate=2e-5, |
|
per_device_train_batch_size=10, |
|
per_device_eval_batch_size=4, |
|
num_train_epochs=3, |
|
weight_decay=0.01, |
|
load_best_model_at_end=True, |
|
metric_for_best_model="accuracy", |
|
logging_dir='logs', |
|
) |
|
|
|
|
|
model = ViTForImageClassification2() |
|
|
|
|
|
def compute_metrics(eval_pred): |
|
predictions, labels = eval_pred |
|
predictions = np.argmax(predictions, axis=1) |
|
return load_metric("accuracy").compute(predictions=predictions, references=labels) |
|
|
|
|
|
trainer = Trainer( |
|
model, |
|
args, |
|
train_dataset=preprocessed_train_ds, |
|
eval_dataset=preprocessed_val_ds, |
|
data_collator=data_collator, |
|
compute_metrics=compute_metrics, |
|
) |
|
|
|
trainer.train() |
|
|
|
outputs = trainer.predict(preprocessed_test_ds) |
