TensorFlow: ¿Cómo calcular el gradiente mientras se ignoran los pesos cero?

Quiero implementar la poda estableciendo algunos pesos en 0.0 (pesos específicos en cada capa, seleccionados de una manera que no es importante discutir aquí) para (1) ignorarlos durante la inferencia y (2) ignorarlos durante el entrenamiento.

Sé que puedo ignorarlos durante la inferencia aplicando una máscara a cada capa, con unas en los índices de ponderaciones que deseo mantener y ceros de lo contrario.

Durante el entrenamiento, sin embargo, quiero que los pesos ignorados permanezcan ignorados. Es decir, quiero que (1) no se actualicen y (2) no afecten el cálculo del rest de los gradientes.

Puedo hacer (1) enmascarando los gradientes antes de aplicarlos. Pero no puedo hacer (2).

He intentado algo como esto:

# Create an optimizer. opt = GradientDescentOptimizer(learning_rate=0.1) # Compute the gradients for a list of variables. # This returns a list of tuples, each tuple containing gradients and # the corresponding variable [(var, grad)] = opt.compute_gradients(loss) # Ask the optimizer to apply the capped gradients. train_op = optimizer.apply_gradients([(grad * mask, var)]) 

Pero esto no hace lo que busco, porque optimizer.compute_gradients(loss) calculará la pérdida en función de todos los pesos, incluidos los que son cero . En otras palabras, grad tendrá algunos gradientes que corresponden a índices donde los pesos son ceros. La aplicación de una máscara “desperdicia” estos gradientes que deberían haberse utilizado para actualizar ponderaciones distintas de cero.

Quiero que calcule la pérdida basada en pesos distintos de cero.

¿Algunas ideas?