Tensorflow multiple optimizers on multi-output model. Get trainable variables for one of the outputs

Question

My goal is to use tfa.optimizers.MultiOptimizer to use a different optimizer for each output of my model. In order to do that I need the layers that feed in to this output, but am unsure how to get those. We can get the model.trainable_variables but this is all the trainable variables and not just those that feed into a given output.

Answer 1

I'm sure there's a better way, but by workaround was to create separate models. Note that using ModelCheckpoint with model.fit will return "not json serializable" error in model.fit. We need to set save_weights_only = True

model = KM.Model(inputs = [in1, in2],outputs=[out1, out2])
model_out1 = KM.Model(inputs = [in1, in2],outputs=[out1])
model_out2 = KM.Model(inputs = [in1, in2],outputs=[out2])
out1_layers = [model.get_layer(j.name) for j in model_out1.layers]
out2_layers = [model.get_layer(j.name) for j in model_out2.layers]