假设您有 3 个输入(例如 roberta 模型 QA 任务)
class MasoudModel2(tf.keras.Model):
def __init__(self):
# in __init__ you define all the layers
super(MasoudModel2, self).__init__()
self.dense1 = tf.keras.layers.Dense(4, activation=tf.nn.relu)
self.dense2 = tf.keras.layers.Dense(10, activation='softmax')
def call(self, inputs):
ids = inputs[0]
toks = inputs[1]
att_mask = inputs[2]
# let's skip real layers
a = self.dense1(ids)
b = self.dense2(att_mask)
return a, b
然后:
ids = tf.keras.Input((MAX_LEN), dtype = tf.int32)
att_mask = tf.keras.Input((MAX_LEN), dtype = tf.int32)
toks = tf.keras.Input((MAX_LEN), dtype = tf.int32)
model2 = MasoudModel2()
model2([ids, att_mask, toks])
更多信息:如果您也想要功能性 API。
def functional_type():
ids = tf.keras.Input((MAX_LEN), dtype = tf.int32)
att_mask = tf.keras.Input((MAX_LEN), dtype = tf.int32)
toks = tf.keras.Input((MAX_LEN), dtype = tf.int32)
c = tf.keras.layers.Dense(10, activation='softmax')(ids)
d = tf.keras.layers.Dense(3, activation = 'softmax')(att_mask)
model = tf.keras.Model(inputs=[ids, toks, att_mask], outputs =[c, d])
return model
然后(注意:第一个参数的最后两个索引是答案。)
model.fit([input_ids[idxT,], attention_mask[idxT,], token_type_ids[idxT,]], [start_tokens[idxT,], end_tokens[idxT,]],
epochs=3, batch_size=32, verbose=DISPLAY, callbacks=[sv],