Just a comment. I’m trying to measure the behaviour of Generative nets, specifically a UNet, so I looked at this, Implemented something in tensorboard and realized that while yes, one does get gradient flow graphs similar to what you show, it isnt quite as straightforwards. The larger the number of weight parameters, the lower the gradient has to be so it does not explode. So you may want to look at the gradients in logscale. Here are 2 representations. The first is similar to the code above, where x:layer number (0 thru 28), y:abs mean gradient (or signed max), z: iteration; uses SummaryWriter.add_histogram_raw() the second x:iteration, y:absmean gradient; using .add_scalars()  
image.jpg757×297 132 KB

Even though the middle gradients are close to zero, they still flow, and the net learns. The higher values are first and last layer which have 5e2 and 1e3 parameters, whereas the parameters of the center layers are~5e6 to 1e7
I tried different inits and the shape of the gradients always ends up being similar, and the UNet learns well. It seems to me that this U shape is a consequence of the netowrk architecture, and shouldnt be necessarily interpreted as wrong.