pytorch bert测试代码中BertForSequenceClassification函数的输入(一条或多条)输出
代码来源是对huggingface的修改，侵删。链接：https://github.com/huggingface/ 现在更新以后大不同了
从后面贴的源码来看，如果是测试，输入就是那三个tensor.1.input_id 2.token_type_ids 3.attention_mask. 训练的话还包括label。
三个特征的代码：参数example是文本的list:[sentence,sentence,sentence....]这里sentence用的是text_a+text_b,没有分句处理。max_seq_length是最大序列长度。tokenizer是 tokenizer = BertTokenizer.from_pretrained(WORK_DIR)得到的。WORK_DIR是放bert模型的vocab.txt的路径，用来加载tokenizer.
def convert_lines(example, max_seq_length, tokenizer):
max_seq_length -= 2
all_tokens = []
all_segments = []
all_masks = []
longer = 0
for text in tqdm(example):
tokens_a = tokenizer.tokenize(text)
if len(tokens_a) > max_seq_length:
tokens_a = tokens_a[:max_seq_length]
longer += 1
one_token = tokenizer.convert_tokens_to_ids(["[CLS]"] + tokens_a + ["[SEP]"]) + [0] * (
max_seq_length - len(tokens_a))
one_segment = [0] * (len(tokens_a) + 2) + [0] * (max_seq_length - len(tokens_a))
one_mask = [1] * (len(tokens_a) + 2) + [0] * (max_seq_length - len(tokens_a))
all_tokens.append(one_token)
all_segments.append(one_segment)
all_masks.append(one_mask)
print(longer)
return np.array(all_tokens), np.array(all_segments), np.array(all_masks)这个代码是有很多条数据的情况，这个函数的输出是三个np类型的数组。all_tokens装所有句子的input_id，是一个n行，max_seq_length列的二维数组。其他两个也是。
部分main函数，
if __name__ == '__main__':
test = pd.read_csv("./dataset/3_abstracts.csv", encoding='utf-8')
test['NAME'] = test['NAME'].fillna("无")
test['CONTENT'] = test['CONTENT'].fillna("无")
test['title_content'] = test['NAME'] + test['CONTENT']
seed_everything()
#######config
device = torch.device('cuda')
WORK_DIR = "./bert_pretrain/"
#我这里分的三类
bert_config = BertConfig.from_pretrained(WORK_DIR + 'bert_config.json', num_labels=3)
tokenizer = BertTokenizer.from_pretrained(WORK_DIR)
MAX_SEQUENCE_LENGTH = 512
test_tokens, test_segments, test_masks = convert_lines(test["title_content"],MAX_SEQUENCE_LENGTH,tokenizer)
#把得到的二维数组包装成tensor.内部为tensor([[id,id...][id..]]),用test_features包装这三个tensor
test_features = [
torch.tensor(test_tokens, dtype=torch.long),
torch.tensor(test_segments, dtype=torch.long),
torch.tensor(test_masks, dtype=torch.long)
]
#pytorch的普遍用法，这个函数把参数处理成一个tensor数据集，是为了后面的loader之类的
test_dataset = torch.utils.data.TensorDataset(*test_features)
#调我的预测函数对标签值进行预测
test_preds = test_model(test_dataset)
预测函数：
def test_model(test_dataset):
WORK_DIR = "./bert_pretrain/"
# WORK_DIR = "./bert_pretrain/"
output_model_file = WORK_DIR + '423_model.bin' #自己训练好的模型
model = BertForSequenceClassification.from_pretrained(WORK_DIR, config=bert_config)
model.load_state_dict(torch.load(output_model_file))
model.to(device)
model.eval()
# for param in model.parameters():
# param.requires_grad = False
test_preds = np.zeros((len(test_dataset), 3))
#SequentialSampler这里是把测试数据顺序排，还有RandomSampler是随机采样，随机排的
test_sampler = SequentialSampler(test_dataset)
#这里加载数据集，主要是设batch是4，也可以设其他，就把刚刚处理过的三兄弟，每个取四个来处理
test_loader = DataLoader(test_dataset, sampler=test_sampler, batch_size=4)
#总的数据量除以4
tk0 = tqdm_notebook(test_loader)
# x_batch1 是一个tentor数据类型：tensor([[id,id..]])，其他两个也是。侧面说明bert的model的输入需要一个二维tensor.源码里有个 num_choices = input_ids.shape[1]这个好像就是求tensor的第二维的长度，我这里设的512
for i, (x_batch1, x_batch2, x_batch3,) in enumerate(tk0):
#注意这里要把数据转到GPU类型，不然会报错
pred = model(x_batch1.to(device), x_batch2.to(device), x_batch3.to(device))
test_preds[i * 4:(i + 1) * 4] = pred[0].detach().cpu().numpy()
return test_preds预测的结果是一个二元组，第二元大概是什么cuda啥啥的，用第一元就行了pred[0],给它转成cpu的numpy。我这里是三分类，得到的结果中[[小数，小数，小数],[],[]...[]]。写一个for循环取三个小数里最大的的索引，就是最终需要的标签。
predict = []
for prediction in test_preds: # predict is one by one, so the length of probabilities=1
pred_label = np.argmax(prediction)
predict.append(pred_label)下面写写如果只输入一条数据有啥要改的，用来部署接口用：
提取特征就只用得到三兄弟了每个是一个list one_token=[id,id,id] one_segment=[seg][seg][seg]...
def convert_lines(example, max_seq_length, tokenizer):
max_seq_length -= 2
longer = 0
tokens_a = tokenizer.tokenize(example)
if len(tokens_a) > max_seq_length:
tokens_a = tokens_a[:max_seq_length]
longer += 1
one_token = tokenizer.convert_tokens_to_ids(["[CLS]"] + tokens_a + ["[SEP]"])+\
[0] * (max_seq_length - len(tokens_a))
one_segment = [0] * (len(tokens_a) + 2) + [0] * (max_seq_length - len(tokens_a))
one_mask = [1] * (len(tokens_a) + 2) + [0] * (max_seq_length - len(tokens_a))
return one_token, one_segment, one_mask因为模型需要一个二维tensor,所以这里转tensor要多加一个中括号。当然也可以unsqueeze(0)
test_token = torch.tensor([test_token])
test_segment = torch.tensor([test_segment])
test_mask=torch.tensor([test_mask])然后测试时这样再这样得到结果，这里我没写全，加载模型啥的跟上面是差不多的
pred = model(test_token.to(device), test_segment.to(device), test_mask.to(device))
predic = pred[0].detach().cpu().numpy()
res = np.argmax(predic)hugging face 的源码：（现在好像更新了封装得更好了）
@add_start_docstrings("""Bert Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. """,
BERT_START_DOCSTRING, BERT_INPUTS_DOCSTRING)
class BertForSequenceClassification(BertPreTrainedModel):
r"""
**labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
Labels for computing the sequence classification/regression loss.
Indices should be in ``[0, ..., config.num_labels - 1]``.
If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).
Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
**loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
Classification (or regression if config.num_labels==1) loss.
**logits**: ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``
Classification (or regression if config.num_labels==1) scores (before SoftMax).
**hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
of shape ``(batch_size, sequence_length, hidden_size)``:
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
**attentions**: (`optional`, returned when ``config.output_attentions=True``)
list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
Examples:
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0) # Batch size 1
labels = torch.tensor([1]).unsqueeze(0) # Batch size 1
outputs = model(input_ids, labels=labels)
loss, logits = outputs[:2]
"""
def __init__(self, config):
super(BertForSequenceClassification, self).__init__(config)
self.num_labels = config.num_labels
self.bert = BertModel(config)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)
self.init_weights()
def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None,
position_ids=None, head_mask=None):
outputs = self.bert(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
attention_mask=attention_mask, head_mask=head_mask)
pooled_output = outputs[1]
pooled_output = self.dropout(pooled_output)
logits = self.classifier(pooled_output)
outputs = (logits,) + outputs[2:] # add hidden states and attention if they are here
if labels is not None:
if self.num_labels == 1:
# We are doing regression
loss_fct = MSELoss()
loss = loss_fct(logits.view(-1), labels.view(-1))
else:
loss_fct = CrossEntropyLoss()
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
outputs = (loss,) + outputs
return outputs # (loss), logits, (hidden_states), (attentions)