I don’t think it’s Python overhead. The GPU utilization as reported by nvidia-smi is relatively high.
Most of the computation is spent in cuDNN convolution. I’m seeing ~4 seconds per epoch which is ~18 ms per iteration (235 iterations per epoch). Convolution forward and backward takes ~13.5 ms PReLU takes ~3 ms (which seems high). The remaining 2ms are probably miscellaneous calls.


  
      gist.github.com
  
  
    https://gist.github.com/colesbury/ea08f4300a52133dfd5201351fcf3b07 5
profile.txt
-------------------------------------  ---------------  ---------------  ---------------  ---------------  ---------------
Name                                          CPU time        CUDA time            Calls        CPU total       CUDA total
-------------------------------------  ---------------  ---------------  ---------------  ---------------  ---------------
N5torch8autograd9GraphRootE                   52.397us          1.408us                1         52.397us          1.408us
expand                                         6.042us          1.584us                2         12.084us          3.168us
fill_                                         10.903us          4.384us                1         10.903us          4.384us
TBackward                                     22.800us          4.224us                2         45.599us          8.448us
view                                           7.472us          4.352us                2         14.945us          8.704us
log_softmax_backward_data                     16.972us          8.800us                1         16.972us          8.800us
ones_like                                     23.823us         10.048us                1         23.823us         10.048us
This file has been truncated. show original




  
  
    
    
  
  


(Note that hierarchy of calls isn’t properly displayed; many of the “convolution” functions are the same call)
What’s the comparable mxnet code you’re running? I think MXNet also relies on cuDNN for convolutions.
EDIT: here’s the script I’m using for profiling: https://gist.github.com/colesbury/2b786128f4409dbd07eef44839f0b3ce 8