pytorch/torch/nn/parallel/parallel_apply.py
Luke Yeager e7c1e6a8e3 [pep8] Fix most lint automatically with autopep8
Here's the command I used to invoke autopep8 (in parallel!):

    git ls-files | grep '\.py$' | xargs -n1 -P`nproc` autopep8 -i

Several rules are ignored in setup.cfg. The goal is to let autopep8
handle everything which it can handle safely, and to disable any rules
which are tricky or controversial to address. We may want to come back
and re-enable some of these rules later, but I'm trying to make this
patch as safe as possible.

Also configures flake8 to match pep8's behavior.

Also configures TravisCI to check the whole project for lint.
2017-01-28 01:15:51 +01:00

48 lines
1.2 KiB
Python

import sys
import threading
import torch
from torch.autograd import Variable
if sys.version_info[0] == 3:
import queue
else:
import Queue as queue
def parallel_apply(modules, inputs):
assert len(modules) == len(inputs)
# Fast track
if len(modules) == 1:
return (modules[0](inputs[0]),)
lock = threading.Lock()
results = {}
def _worker(module, input, results, lock):
var_input = input
while not isinstance(var_input, Variable):
var_input = var_input[0]
try:
with torch.cuda.device_of(var_input):
output = module(input)
with lock:
results[input] = output
except Exception as e:
with lock:
results[input] = e
threads = [threading.Thread(target=_worker,
args=(module, input, results, lock))
for module, input in zip(modules, inputs)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
outputs = []
for i in inputs:
output = results[i]
if isinstance(output, Exception):
raise output
outputs.append(output)
return outputs