| 629 | |
| 630 | template <typename Gradient, typename BackwardFunction, typename TapeTensor> |
| 631 | Status GradientTape<Gradient, BackwardFunction, TapeTensor>::ComputeGradient( |
| 632 | const VSpace<Gradient, BackwardFunction, TapeTensor>& vspace, |
| 633 | const gtl::ArraySlice<int64> target_tensor_ids, |
| 634 | const gtl::ArraySlice<int64> source_tensor_ids, |
| 635 | const std::unordered_map<int64, TapeTensor>& sources_that_are_targets, |
| 636 | gtl::ArraySlice<Gradient*> output_gradients, |
| 637 | std::vector<Gradient*>* result) { |
| 638 | std::unordered_set<int64> sources_set(source_tensor_ids.begin(), |
| 639 | source_tensor_ids.end()); |
| 640 | BackpropInitialState<BackwardFunction, TapeTensor> state = PrepareBackprop( |
| 641 | target_tensor_ids, tensor_tape_, &op_tape_, sources_set, persistent_); |
| 642 | std::vector<int64> op_stack = |
| 643 | InitialStack(state.op_tape, state.op_missing_tensor); |
| 644 | std::unordered_map<int64, std::vector<Gradient*>> gradients; |
| 645 | Status s = InitialGradients(vspace, target_tensor_ids, |
| 646 | sources_that_are_targets, output_gradients, |
| 647 | tensor_tape_, state.op_tape, &gradients); |
| 648 | auto cleanup = [this, &state]() { |
| 649 | if (!persistent_) { |
| 650 | // Release all backprop functions |
| 651 | for (const auto& pair : state.op_tape) { |
| 652 | pair.second.backward_function_deleter(pair.second.backward_function); |
| 653 | } |
| 654 | } |
| 655 | }; |
| 656 | if (!s.ok()) { |
| 657 | cleanup(); |
| 658 | return s; |
| 659 | } |
| 660 | |
| 661 | std::unordered_map<int64, int64> gradients_size; |
| 662 | // TODO(apassos) multiple threads could be dequeuing from op_stack at the same |
| 663 | // time, for better CPU backprop performance. |
| 664 | VLOG(1) << "Initial stack:"; |
| 665 | if (VLOG_IS_ON(1)) { |
| 666 | for (auto t : op_stack) { |
| 667 | VLOG(1) << " " << t; |
| 668 | } |
| 669 | } |
| 670 | while (!op_stack.empty()) { |
| 671 | const int64 op = op_stack.back(); |
| 672 | VLOG(1) << "Popped " << op; |
| 673 | op_stack.pop_back(); |
| 674 | auto op_it = state.op_tape.find(op); |
| 675 | if (op_it == state.op_tape.end()) { |
| 676 | // It is possible for ops to end up on the stack if they are unrelated to |
| 677 | // the target; we should just skip them. |
| 678 | continue; |
| 679 | } |
| 680 | auto trace = std::move(op_it->second); |
| 681 | state.op_tape.erase(op_it); |
| 682 | std::vector<Gradient*> out_gradients; |
| 683 | out_gradients.reserve(trace.output_tensor_info.size()); |
| 684 | std::vector<int64> unneeded_gradients; |
| 685 | for (size_t i = 0; i < trace.input_tensor_id.size(); i++) { |
| 686 | const auto& in_tensor_id = trace.input_tensor_id[i]; |
| 687 | if (tensor_tape_.find(in_tensor_id) == tensor_tape_.end() && |
| 688 | sources_set.find(in_tensor_id) == sources_set.end()) { |
no test coverage detected