This is a nice start, but what would really help is for someone who understands the GPU programming model in-depth to give this a shot, with the goal of reducing DRAM bandwidth and fitting the layers exactly onto a GPU's memory hierarchy (cache levels, local memory and registers). Basically, sizes of a target HW platform's memories should be hyper-parameters of the architecture.