(proto, snapshot, gpus, timing, uid, rank)
| 63 | |
| 64 | |
| 65 | def solve(proto, snapshot, gpus, timing, uid, rank): |
| 66 | caffe.set_mode_gpu() |
| 67 | caffe.set_device(gpus[rank]) |
| 68 | caffe.set_solver_count(len(gpus)) |
| 69 | caffe.set_solver_rank(rank) |
| 70 | caffe.set_multiprocess(True) |
| 71 | |
| 72 | solver = caffe.SGDSolver(proto) |
| 73 | if snapshot and len(snapshot) != 0: |
| 74 | solver.restore(snapshot) |
| 75 | |
| 76 | nccl = caffe.NCCL(solver, uid) |
| 77 | nccl.bcast() |
| 78 | |
| 79 | if timing and rank == 0: |
| 80 | time(solver, nccl) |
| 81 | else: |
| 82 | solver.add_callback(nccl) |
| 83 | |
| 84 | if solver.param.layer_wise_reduce: |
| 85 | solver.net.after_backward(nccl) |
| 86 | solver.step(solver.param.max_iter) |
| 87 | |
| 88 | |
| 89 | if __name__ == '__main__': |
nothing calls this directly
no test coverage detected