Tensorflow中如何对模型简单剪枝？

def getZeroIndex(array):
    res=[]
    for index,value in enumerate(array):
        if value==0:
            res.append(index)
    return np.array(res)

watermark_model=tf.keras.models.load_model("hn_teacher_watermark_97.h5")
watermark_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
            metrics=['accuracy'])
for i in range(len(watermark_model.layers)):
    t=watermark_model.layers[i].get_weights()
    if len(t)==2:
        print(watermark_model.layers[i].name)
        weight=t[0]
        origin_shape=weight.shape
        weight=weight.flatten()
        change_weight=weight*weight
        change_weight[change_weight<0.0004]=0
        index=getZeroIndex(change_weight)
        print(index.shape)
        weight[index]=0
        weight=weight.reshape(origin_shape)
        t[0]=weight
        watermark_model.layers[i].set_weights(t)
watermark_model.evaluate(data.teacher_clean_test_x,data.teacher_clean_test_y)
print("{:.4f}".format(WatermarkDetectiveRateInTrain(watermark_model))+"/"+"{:.4f}".format(WatermarkDetectiveRate(watermark_model)))