系统内存崩溃和CUDA内存不足问题是在使用BERT模型进行分类任务时常见的问题。下面是一些解决这些问题的方法,包括代码示例:
from transformers import BertTokenizer, BertForSequenceClassification
import torch
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
# 输入数据和标签
data = ["This is the first sentence.", "This is the second sentence."]
labels = [0, 1]
# 分词和编码
encoded_inputs = tokenizer(data, padding=True, truncation=True, return_tensors='pt')
labels = torch.tensor(labels)
# 设置较小的batch size
batch_size = 2
# 按batch size进行迭代
for i in range(0, len(data), batch_size):
batch_inputs = {k: v[i:i+batch_size] for k, v in encoded_inputs.items()}
batch_labels = labels[i:i+batch_size]
# 将输入传递给模型
outputs = model(**batch_inputs, labels=batch_labels)
loss = outputs.loss
logits = outputs.logits
# 后续处理...
from transformers import BertTokenizer, BertForSequenceClassification
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased').to(device)
# 输入数据和标签
data = ["This is the first sentence.", "This is the second sentence."]
labels = [0, 1]
# 分词和编码
encoded_inputs = tokenizer(data, padding=True, truncation=True, return_tensors='pt').to(device)
labels = torch.tensor(labels).to(device)
# 设置较小的batch size
batch_size = 2
# 按batch size进行迭代
for i in range(0, len(data), batch_size):
batch_inputs = {k: v[i:i+batch_size] for k, v in encoded_inputs.items()}
batch_labels = labels[i:i+batch_size]
# 将输入传递给模型
outputs = model(**batch_inputs, labels=batch_labels)
loss = outputs.loss
logits = outputs.logits
# 后续处理...
这些方法可以帮助你解决BERT模型分类时的系统内存崩溃和CUDA内存不足问题。根据你的具体情况,你可以选择其中一个或两个方法来解决问题。