给定一个三维Keras张量，如何将尺寸1和2的对角线设置为1？

给定Keras张量，我想用数字1代替对角线。这类似于numpy函数

fill_diagonal

让

inputTensor

成为您要更改的张量，

size

成为尺寸1和2的大小：

#create a diagonal matrix
diag = K.eye(size)

#make it 3D by adding a dummy batch dimension
diag = K.expand_dims(diag,0)

#get a complement for wich complement + diag = ones    
complement = 1 - diag

#annulate the diagonal from the input
noDiagInput = complement * inputTensor    

#add diagonal 1
return noDiagInput + diag