(solver, nccl)
| 31 | |
| 32 | |
| 33 | def time(solver, nccl): |
| 34 | fprop = [] |
| 35 | bprop = [] |
| 36 | total = caffe.Timer() |
| 37 | allrd = caffe.Timer() |
| 38 | for _ in range(len(solver.net.layers)): |
| 39 | fprop.append(caffe.Timer()) |
| 40 | bprop.append(caffe.Timer()) |
| 41 | display = solver.param.display |
| 42 | |
| 43 | def show_time(): |
| 44 | if solver.iter % display == 0: |
| 45 | s = '\n' |
| 46 | for i in range(len(solver.net.layers)): |
| 47 | s += 'forw %3d %8s ' % (i, solver.net._layer_names[i]) |
| 48 | s += ': %.2f\n' % fprop[i].ms |
| 49 | for i in range(len(solver.net.layers) - 1, -1, -1): |
| 50 | s += 'back %3d %8s ' % (i, solver.net._layer_names[i]) |
| 51 | s += ': %.2f\n' % bprop[i].ms |
| 52 | s += 'solver total: %.2f\n' % total.ms |
| 53 | s += 'allreduce: %.2f\n' % allrd.ms |
| 54 | caffe.log(s) |
| 55 | |
| 56 | solver.net.before_forward(lambda layer: fprop[layer].start()) |
| 57 | solver.net.after_forward(lambda layer: fprop[layer].stop()) |
| 58 | solver.net.before_backward(lambda layer: bprop[layer].start()) |
| 59 | solver.net.after_backward(lambda layer: bprop[layer].stop()) |
| 60 | solver.add_callback(lambda: total.start(), lambda: (total.stop(), allrd.start())) |
| 61 | solver.add_callback(nccl) |
| 62 | solver.add_callback(lambda: '', lambda: (allrd.stop(), show_time())) |
| 63 | |
| 64 | |
| 65 | def solve(proto, snapshot, gpus, timing, uid, rank): |
no test coverage detected