Open LIU-FAYANG opened 2 years ago
AFAIK evenif you apply prune_low_magnitude to sublayers of your subclass model, the pruning scheduling logic (callback based) during the training is not working due to it can't find your pruned layer. So I think your model is not trained as expected with pruning.
Unfortunately, subclass model supports is very weak in these days. Maybe there's some work around but it's still a brittle solution.
@rino20 Hi rino20, do you know any recent workaround for the subclass model supports? or any recommendation?
Hi @Xhark, thanks for your kind help. This page https://github.com/tensorflow/model-optimization/issues/155 mentioned apply the pruning API directly into subclasses is a workaround method. Currently I'm trying to prune a RNN-T model using prune_low_magnitude and the pruning result is indeed not good, the model tends not to converge and I'm not sure if it's because of this workaround method. Could u elaborate more about why the pruning scheduling logic is not working in this way? I tired to make to all of the layers to be pruned have the same pruning schedule so I'm a bit confused why the pruning scheduling logic is not working this way. Thanks for your kind help!
Hi @LIU-FAYANG
I am not sure how you apply the pruning to subclass model, but if you applied pruning for the keras Layers in the subclass models, it should be working.
The general strip_pruning method wouldn't work in that case, since you applied pruning in custom way. You will need to create your own strip_pruning method which is the same way of your applying pruning logic.
Could you elaborate more about "model is not converging"? If model is alternating a lot, longer pruning_frequency is recommended.
Hi @rino20, thanks for your help!
Thanks for your kind help!
Hi Thanks for providing the details.
Yes, the way you took is working. In detail, you can achieve your goal by 1) prune (pruning_wrapping) the target layer, and 2) build the subclass model.
Your strip_pruning should do the same as #1. it should 1) get the pruning_wrapped layer -> apply strip_pruning 2) build the subclass model Or you can just swap the layer with stripped layer, however, note that this is not the best practice. run_transducer_block.projection = tfmot.strip_pruning(run_transducer_block.projection)
The loss exploding even with begin_step != 0 seems weird. Could you test with 0 sparsity level and inspect the actual weight value during training?When you proceed training, the loss reduced?
Hi @rino20, thanks for your kind help, I'll try the method u mentioned above these days. Another thing I would like to ask is the model I'm working on seems have custom training loop, I didn't merge the pruning callbacks to the custom training loop but the model is still able to train, which is quite weird. I'd like to ask is it necessary to merge the callbacks to the training loop in this case? Or should I just add the callback list to the model.fit() like normal training? Maybe this is what Xhark means when he mentioned the concern of pruning callbacks for subclass models?
edit: I tired on a small end to end subclass model used in https://www.tensorflow.org/guide/keras/custom_layers_and_models, which has custom training loop for this particular model, I tired to add the pruning callback to the custom training loop by referring to the pruning comprehensive guide but the mask & weights seems does not change at all after pruning for a few epochs. May I ask is it because the pruning callback does not support subclass model?
Hi, is there any work around method to use tfmot.sparsity.keras.strip_pruning on subclass models? I tried to use prune_low_magnitude on subclass by apply it to the layers in subclasses as a work around method. So I would like to ask is there method like this to use strip_pruning on subclasses models? Any help would be great:)