#pragma once #include "caffe2/core/net.h" #include "caffe2/core/observer.h" #include "caffe2/core/operator.h" namespace caffe2 { // Thin class that attaches the observer to all operators in the net template class OperatorAttachingNetObserver : public ObserverBase { public: explicit OperatorAttachingNetObserver( NetBase* subject_, TNetObserver* netObserver) : ObserverBase(subject_) { const auto& operators = subject_->GetOperators(); for (auto* op : operators) { auto observer = caffe2::make_unique(op, netObserver); const auto* ob = observer.get(); op->AttachObserver(std::move(observer)); operator_observers_.push_back(ob); } } virtual ~OperatorAttachingNetObserver(){}; protected: std::vector operator_observers_; }; } // namespace caffe2