GLUE_TASKS = ["cola", "mnli", "mnli-mm", "mrpc", "qnli", "qqp", "rte", "sst2", "stsb", "wnli"]
#useful in preprocessing, this sets what each task does.
task_to_keys = {
"cola": ("sentence", None),
"mnli": ("premise", "hypothesis"),
"mnli-mm": ("premise", "hypothesis"),
"mrpc": ("sentence1", "sentence2"),
"qnli": ("question", "sentence"),
"qqp": ("question1", "question2"),
"rte": ("sentence1", "sentence2"),
"sst2": ("sentence", None),
"stsb": ("sentence1", "sentence2"),
"wnli": ("sentence1", "sentence2"),
}
def set_module_grad_status(module, flag=False):
if isinstance(module, list):
for m in module:
set_module_grad_status(m, flag)
else:
for p in module.parameters():
p.requires_grad = flag
def compute_metrics(eval_pred):
predictions, labels = eval_pred
if task != "stsb":
predictions = np.argmax(predictions, axis=1)
else:
predictions = predictions[:, 0]
return metric.compute(predictions=predictions, references=labels)
#setting some model parameters
task = "qnli"
sentence1_key, sentence2_key = task_to_keys[task]
# model_checkpoint = "test-glue/basemodel"
num_labels = 3 if task.startswith("mnli") else 1 if task=="stsb" else 2
metric_name = "pearson" if task == "stsb" else "matthews_correlation" if task == "cola" else "accuracy"
validation_key = "validation_mismatched" if task == "mnli-mm" else "validation_matched" if task == "mnli" else "validation"
model_checkpoint = "google/mobilebert-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint, num_labels=num_labels)
model = model.to(device)
args = TrainingArguments(
"test-glue",
evaluation_strategy = "epoch",
learning_rate=2e-5,
per_device_train_batch_size=batch_size,
per_device_eval_batch_size=batch_size,
num_train_epochs=epochs,
weight_decay=0.01,
save_steps=0
# load_best_model_at_end=True,
# metric_for_best_model=metric_name,
)
#load dataset
actual_task = "mnli" if task == "mnli-mm" else task #some error checking..
dataset = load_dataset("glue", actual_task)
metric = load_metric('glue', actual_task)
# tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
def preprocess_function(examples):
if sentence2_key is None:
return tokenizer(examples[sentence1_key], truncation=True)
return tokenizer(examples[sentence1_key], examples[sentence2_key], truncation=True)
encoded_dataset = dataset.map(preprocess_function, batched=True)
trainer = Trainer(
model,
args,
train_dataset=encoded_dataset["train"],
eval_dataset=encoded_dataset[validation_key],
tokenizer=tokenizer,
compute_metrics=compute_metrics
)
trainer.train()
我不太清楚为什么会收到错误“RuntimeError:张量 a (549) 的大小必须与非单维 1 处的张量 b (512) 的大小相匹配”,例如“trainer.train()”。就像不是每个样本都应该是一致的 512 大小。这只发生在 MobileBERT 模型上,而不是在 DistilBERT 模型上。我使用的代码主要基于 huggingface colab 教程https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification.ipynb#scrollTo=7k8ge1L1IrJk