File size: 437 Bytes
5fa1a76 |
1 2 3 4 5 6 7 8 9 10 11 12 13 |
Specify the number of labels along with the number of expected labels, and the label mappings:
from transformers import AutoModelForImageClassification, TrainingArguments, Trainer
model = AutoModelForImageClassification.from_pretrained(
checkpoint,
num_labels=len(labels),
id2label=id2label,
label2id=label2id,
)
At this point, only three steps remain:
Define your training hyperparameters in [TrainingArguments]. |