| from collections import deque, defaultdict |
| from torch._C import _ImperativeEngine as ImperativeEngine |
| from .variable import Variable |
| |
| |
| class BasicEngine(object): |
| |
| def _compute_dependencies(self, function): |
| dependencies = defaultdict(int) |
| seen = {function} |
| queue = [function] |
| while len(queue) > 0: |
| fn = queue.pop() |
| for prev_fn, output_nr in fn.previous_functions: |
| if not prev_fn.requires_grad or isinstance(prev_fn, Variable): |
| continue |
| dependencies[prev_fn] += 1 |
| if prev_fn not in seen: |
| queue.append(prev_fn) |
| seen.add(prev_fn) |
| return dependencies |
| |
| def _free_backward_dependency(self, dependencies, prev_fn): |
| dependencies[prev_fn] -= 1 |
| if dependencies[prev_fn] == 0: |
| del dependencies[prev_fn] |
| return True |
| return False |
| |
| def _add_grad(self, need_copy, prev_grad, output_nr, d_prev_fn): |
| copy_id = (id(prev_grad), output_nr) |
| if not prev_grad[output_nr]: |
| prev_grad[output_nr] = d_prev_fn |
| need_copy.add(copy_id) |
| else: |
| grad_tensor = prev_grad[output_nr] |
| if copy_id in need_copy: |
| need_copy.remove(copy_id) |
| grad_tensor = grad_tensor.clone() |
| prev_grad[output_nr] = grad_tensor |
| grad_tensor.add_(d_prev_fn) |
| |
| def run_backward(self, variable, grad, retain_variables): |
| if variable.creator is None: |
| variable._do_backward((grad,), retain_variables) |
| return |
| |
| initial_grad = [None for _ in range(variable.creator.num_outputs)] |
| initial_grad[variable.output_nr] = grad |
| ready = deque([(variable.creator, initial_grad)]) |
| not_ready = {} |
| need_copy = set() |
| |
| dependencies = self._compute_dependencies(variable.creator) |
| |
| while len(ready) > 0: |
| fn, grad = ready.pop() |
| grad_input = fn._do_backward(tuple(grad), retain_variables) |
| for (prev_fn, output_nr), d_prev_fn in zip(fn.previous_functions, grad_input): |
| if not prev_fn.requires_grad: |
| # TODO: check that d_prev_fn is None and warn otherwise |
| continue |
| if isinstance(prev_fn, Variable): |
| prev_fn._do_backward((d_prev_fn,), retain_variables) |
| continue |
| is_ready = self._free_backward_dependency(dependencies, prev_fn) |
| if is_ready: |
| if prev_fn in not_ready: |
| prev_grad = not_ready[prev_fn] |
| self._add_grad(need_copy, prev_grad, output_nr, d_prev_fn) |
| else: |
| if prev_fn.num_outputs != 1: |
| raise RuntimeError("one of the function outputs " |
| "wasn't used - this is an error not, but " |
| "it's going to be fixed soon") |
| prev_grad = (d_prev_fn,) |
| ready.appendleft((prev_fn, prev_grad)) |
| else: |
| if prev_fn in not_ready: |
| prev_grad = not_ready[prev_fn] |
| else: |
| prev_grad = [None for _ in range(prev_fn.num_outputs)] |
| |
| self._add_grad(need_copy, prev_grad, output_nr, d_prev_fn) |
| not_ready[prev_fn] = prev_grad |
