| 598 | } |
| 599 | |
| 600 | void LoopGrad::GradDesc::init_virtual_loss( |
| 601 | DescImplBase* fwd_desc, const VarNodeArray& outgrad_owner) { |
| 602 | GradProxy::GradInfoArray output_grad_info; |
| 603 | VarNodeArray needed_fwd_outvars; |
| 604 | |
| 605 | // handle user added outputs: forward grads in owner graph to subgraph |
| 606 | size_t idx = 0; |
| 607 | |
| 608 | // user_data records whether it has been added |
| 609 | for (auto&& i : fwd_desc->output_record_spec_no_dedup()) |
| 610 | i->user_data = false; |
| 611 | for (auto&& i : fwd_desc->output_record_spec_no_dedup()) { |
| 612 | if (!i->user_data) { |
| 613 | i->user_data = true; |
| 614 | auto owner_grad = outgrad_owner.at(idx); |
| 615 | if (owner_grad) { |
| 616 | auto rec = i->recorder(); |
| 617 | auto all_grad_sub = add_input(owner_grad); |
| 618 | auto sub_grad = rec->get_outgrad_in_iter( |
| 619 | get_counter_var(), m_counter_var_up, all_grad_sub); |
| 620 | sub_grad.rename(ssprintf("outgrad:%s[%zd]", rec->name().c_str(), idx)); |
| 621 | needed_fwd_outvars.push_back(i->var_sub()); |
| 622 | output_grad_info.push_back({i->var_sub(), sub_grad.node()}); |
| 623 | } |
| 624 | idx++; |
| 625 | } |
| 626 | } |
| 627 | |
| 628 | mgb_assert(idx == outgrad_owner.size()); |
| 629 | m_fwd_graph_modifier.init(needed_fwd_outvars); |
| 630 | for (auto&& i : output_grad_info) { |
| 631 | i.wrt = m_fwd_graph_modifier.map_var(i.wrt); |
| 632 | } |
| 633 | |
| 634 | m_grad_virtual_loss = |
| 635 | GradProxy::make(m_sub_graph.get(), std::move(output_grad_info)); |
| 636 | m_grad_virtual_loss_opr = |
| 637 | &m_grad_virtual_loss.node()->owner_opr()->cast_final_safe<GradProxy>(); |
| 638 | } |
| 639 | |
| 640 | bool LoopGrad::GradDesc::bind_grad_var(VarNode* owner_wrt, VarNode* owner_dest) { |
| 641 | auto&& input_ogvar2info = m_fwd_graph_modifier.input_ogvar2info(); |
nothing calls this directly
no test coverage detected