Updates gradient_generators and gradient_frontier
( # NOQA
self, fwd_op_idx, gradient_ops, g_output, g_input)
| 652 | self.gradient_generators[name][version].append(generator) |
| 653 | |
| 654 | def BuildGradientGenerators( # NOQA |
| 655 | self, fwd_op_idx, gradient_ops, g_output, g_input): |
| 656 | """Updates gradient_generators and gradient_frontier""" |
| 657 | forward_op, in_versions, out_versions = self.ssa[fwd_op_idx] |
| 658 | locally_generated_blobs = [] |
| 659 | sparse_generators = defaultdict(lambda: defaultdict(list)) |
| 660 | |
| 661 | for grad_op in gradient_ops: |
| 662 | # (1) check that inputs are valid |
| 663 | for s in grad_op.input: |
| 664 | self.CheckGradientOperatorInput( |
| 665 | s, g_output, fwd_op_idx, locally_generated_blobs) |
| 666 | |
| 667 | # (2) add outputs to the locally generated blobs |
| 668 | # If an output corresponds to the gradient of an input, we also |
| 669 | # record it to gradient_generators |
| 670 | locally_generated_blobs.extend(map(str, grad_op.output)) |
| 671 | for i, output in enumerate(grad_op.output): |
| 672 | input_index = GetIndexFromGradientList(g_input, output) |
| 673 | if input_index is not None: |
| 674 | input_name = forward_op.input[input_index] |
| 675 | input_version = in_versions[input_name] |
| 676 | g = g_input[input_index] |
| 677 | if type(g) is GradientSlice: |
| 678 | # the output corresponds either to the indices or the |
| 679 | # values of the sparse gradient. In either case we |
| 680 | # create a (partial) SparseGradGenMeta. If necessary, |
| 681 | # we'll merge indices and values generators |
| 682 | # corresponding to the same gradient in step (3) |
| 683 | if g.indices == output: |
| 684 | m = SparseGradGenMeta( |
| 685 | grad_op, i, None, 0, g, grad_op.device_option) |
| 686 | else: |
| 687 | assert(g.values == output) |
| 688 | m = SparseGradGenMeta( |
| 689 | None, 0, grad_op, i, g, grad_op.device_option) |
| 690 | sparse_generators[input_name][input_version].append(m) |
| 691 | else: |
| 692 | self.gradient_generators[input_name][input_version] \ |
| 693 | .append(GradGenMeta( |
| 694 | grad_op, i, g, grad_op.device_option)) |
| 695 | |
| 696 | # (3) merge indices and values generators for sparse gradients, and |
| 697 | # add them to gradient_generators |
| 698 | self.AppendSparseGenerators(sparse_generators) |
| 699 | |
| 700 | # (4) for ops (e.g., Add, Sum, Sub) which have gradient outputs directly |
| 701 | # passed from inputs (not computed from gradient ops), we create an |
| 702 | # GradGenMeta with None grad_op and idx so that the gradient_generators |
| 703 | # knows where the gradients are coming from. This is needed for creating |
| 704 | # Sum op to accumulate the gradients from multiple parents. |
| 705 | for input_index, g in enumerate(g_input): |
| 706 | input_name = forward_op.input[input_index] |
| 707 | input_version = in_versions[input_name] |
| 708 | if not g: |
| 709 | continue |
| 710 | if type(g) is GradientSlice: |
| 711 | if str(g.indices) not in locally_generated_blobs and \ |
no test coverage detected