使用TFMOT进行TensorFlow模型剪枝时训练验证损失均为nan如何解决？

你当前的问题核心是错误地将自定义Instance Normalization层的所有权重标记为可剪枝，且对全模型无差别应用剪枝策略，破坏了模型的数值稳定性，进而导致损失为NaN、输出全0，可按以下步骤排查解决：

• 第一步：修改自定义instance_normalization层的剪枝逻辑
  归一化层参数规模极小，剪枝收益极低，且对其权重剪枝极易导致归一化计算出现除0、数值溢出等问题。直接将该层的get_prunable_weights修改为返回空列表，跳过该层的剪枝：

  ### defined for model pruning
  def get_prunable_weights(self):
      return []
  

  如果你确实需要对该层的缩放参数gamma做剪枝，也仅返回gamma对应的权重，不要返回beta、滑动均值、滑动方差等非剪枝目标参数。
• 第二步：限定剪枝的应用范围，避免对非必要层剪枝
  你当前直接对全模型调用prune_low_magnitude，会对池化、激活、归一化等无参数/不适合剪枝的层都套上剪枝 wrapper，进一步提升了数值不稳定的概率。改为仅对卷积层、全连接层等参数密集的层应用剪枝：

  from tensorflow_model_optimization.sparsity.keras import prune_low_magnitude
  
  def apply_pruning_to_layer(layer):
      # 仅对Conv2D、Dense层应用剪枝，可根据你的模型结构调整适配
      if isinstance(layer, (tf.keras.layers.Conv2D, tf.keras.layers.Dense)):
          return prune_low_magnitude(layer, **pruning_params)
      return layer
  
  # 替换原有的全模型prune_low_magnitude调用
  model_for_pruning = tf.keras.models.clone_model(
      model,
      clone_function=apply_pruning_to_layer
  )
  

• 第三步：调整剪枝调度策略，降低初始扰动
  你当前配置的begin_step=0、initial_sparsity=0.5会在训练一开始就直接将50%的权重置0，对预训练好的基础模型权重分布破坏过大。可以调整为逐步启动剪枝：

  pruning_params = {
      'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
          initial_sparsity=0.0, # 初始不剪枝
          final_sparsity=0.80,
          begin_step=100, # 训练100步后再开始剪枝，给模型适配时间
          end_step=end_step)
  }
  

• 第四步：校验损失函数配置
  确认你的基础模型输出层没有加Softmax激活，如果基础模型输出层已经添加了Softmax，需要将损失函数的from_logits=True改为from_logits=False，避免交叉熵计算出现数值溢出。
• 第五步：添加梯度裁剪避免梯度爆炸
  编译模型时给优化器添加梯度裁剪，防止剪枝过程中梯度过大冲毁参数：

  model_for_pruning.compile(
      optimizer=tf.keras.optimizers.Adam(clipnorm=1.0),
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=['accuracy']
  )
  

内容的提问来源于stack exchange，提问作者E_learner