' Dense(units=1) '输出形状不是' (batch_size, 1) '

Question

这是一个模型:

import random
import string
import tensorflow as tf
from keras import Input, Model
from keras.layers import Dense, Dropout
from transformers import DistilBertTokenizer, TFDistilBertModel

def create_model(input_shape, cache_dir, pretrained_model):
ids = Input(input_shape, dtype='int32', name='input_ids')
masks = Input(input_shape, dtype='int32', name='attention_mask')
x0 = dict(input_ids=ids, attention_mask=masks)
x = TFDistilBertModel.from_pretrained(pretrained_model, cache_dir=cache_dir)(x0)[0]
x = Dropout(0.2)(x)
x = Dense(64)(x)
x = Dropout(0.2)(x)
output = Dense(1, 'sigmoid')(x)
return Model(x0, output))

我需要输出形状为(100, 1)，然而，在这个例子中它是(100, 512):

pretrained_model = 'distilbert-base-uncased'
cache_dir = '.'
tokenizer = DistilBertTokenizer.from_pretrained(
pretrained_model, cache_dir=cache_dir
)
s = string.ascii_letters
text = [''.join(random.choice(s) for _ in range(10)) for _ in range(100)]
tokens = tokenizer.batch_encode_plus(
text,
add_special_tokens=True,
truncation=True,
return_tensors='np',
padding='max_length',
)
model = create_model((None,), cache_dir, pretrained_model)
preds = model(dict(tokens))
print(preds.shape)

:

TensorShape([100, 512])

这意味着获得n个标记化的句子作为输入，并为每个句子输出一个分数，以确定它们的顺序。如何使用代码的当前/其他变体来实现这一点?

Answer 1

除非我错过了什么，否则您将返回具有概率的完整向量。如果是这种情况，

np.argmax(preds,axis=1)

…可以解决你的问题。