| 2204 | |
| 2205 | @staticmethod |
| 2206 | def _RunAllOnGPU(net, gpu_id=0, use_cudnn=False): |
| 2207 | device_option = caffe2_pb2.DeviceOption() |
| 2208 | device_option.device_type = workspace.GpuDeviceType |
| 2209 | device_option.device_id = gpu_id |
| 2210 | net.device_option.CopyFrom(device_option) |
| 2211 | if use_cudnn: |
| 2212 | for op in net.op: |
| 2213 | op.engine = "CUDNN" |
| 2214 | # Move RecurrentNetwork operators on GPU as well |
| 2215 | for op in net.op: |
| 2216 | if op.type != "RecurrentNetwork": |
| 2217 | continue |
| 2218 | for arg in op.arg: |
| 2219 | if arg.name == "step_net": |
| 2220 | Net._RunAllOnGPU(arg.n, gpu_id, use_cudnn) |
| 2221 | |
| 2222 | def RunAllOnGPU(self, gpu_id=0, use_cudnn=False): |
| 2223 | """A convenient function to run everything on the GPU.""" |