【tf.keras】：版本从1.x升级到2.x遇到的一个问题：InvalidArgumentError: Cannot assign a device for operation embedding_

问题

【tf.keras】：版本从1.x升级到2.x遇到的一个问题：

tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device for operation model_1/Embedding-Norm/mul/ReadVariableOp: Could not satisfy explicit device specification '' because the node {
    {
    colocation_node model_1/Embedding-Norm/mul/ReadVariableOp}} was colocated with a group of nodes that required incompatible device '/job:localhost/replica:0/task:0/device:GPU:0'. All available devices [/job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:GPU:0]. 

解决

1.x的优化函数是这样的：

model = Model(model.inputs, output)
model.summary()

AdamW = extend_with_weight_decay(Adam, 'AdamW')
AdamWG = extend_with_gradient_accumulation(AdamW, 'AdamWG')
optimizer = AdamWG(
    learning_rate=2e-5,
    weight_decay_rate=0.01,
    exclude_from_weight_decay=['Norm', 'bias'],
    grad_accum_steps=16
)
model.compile(optimizer=optimizer)

2.x的优化函数需要改成这样的：

model = Model(model.inputs, output)
model.summary()

AdamW = extend_with_weight_decay(Adam, 'AdamW')
optimizer = AdamW(learning_rate=2e-6, weight_decay_rate=0.01)
model.compile(optimizer=optimizer)