CODEWITHSUNDEEP
tensorflowkeras
AstroBen
AstroBen

Reputation: 863

Where do you put 'mixed_precision' policy in Keras models in a large library with many files?

Using Tensorflow 2.3, I have a keras model with function definitions that span over a dozen files. There is one master file that runs the whole thing and does the fitting, but of course each file has its own import statements. The model is built in build.py, compiled in compile.py, and then run from the master.py. I know that if I want to train in mixed precision, I need to execute the following (or some longer/shorter variation) before compilation:

from tensorflow.keras.mixed_precision import experimental as mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

My question is, do I need to declare this in only the master.py file, the file where the model is defined (build.py), compiled (compiled.py), or all files that have anything to do with defining the model?

Upvotes: 1

Views: 228

Answers (1)

Richard X
Richard X

Reputation: 1134

You will want to set the mixed precision policy inside the file where the model is being built, which in this case is build.py. The reason being that the data type of a layer will be inferred from the global policy by default. Note that this happens when the layer is being constructed. If your build.py contains a function which returns a model, then you can do as before or set the global policy in the file where you call the function to construct the model (master.py).

Upvotes: 1

Related Questions