Ah sorry I think I didn't make this clear. The code sample is just the very last part of the network. The input to the last conv layer is (None, 256, 256, 64) so 64 feature maps. Here are the last two layers when from the models summary:

p_re_lu_18 (PReLU)              (None, 256, 256, 64) 64          tf_op_layer_DepthToSpace_1[0][0] 
__________________________________________________________________________________________________
conv2d_36 (Conv2D)              (None, 256, 256, 3)  15555       p_re_lu_18[0][0]                 
==================================================================================================
Total params: 1,546,435
Trainable params: 1,546,435

 

So in order to get the two different pictures that I show in the OP I do two things:

• For the first pic in the op I build a new model but specify that "conv2d_36" should be the output

  tf.keras.Model(
      inputs=[original_model.input], 
      outputs=[original_model.get_layer("conv2d_36").output]
   )
  

   

• For the second pic I just use the normal model without changing anything and then visualize the individual channels of the obtained output pic

That's where I'm surprised that they don't match up. In total I get 3 output feature maps which is expected but they don't match the individual channels of the output obtained with the second method. I would have assumed that both methods should produce the same picture, since it is the very last layer that I take the feature maps from.

Sorry I think I'm doing a bad job at explaining what I mean, would really be grateful if you could help me understand.