mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
48db74ea03
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/13370 This diff adds a new net type (simple_refcount) that does one thing: for all intermediate results produced by a net, it will keep refcount about internal usage, and when it finishes its consumption, the net will delete the blob content to mimic the case of dynamic allocation. In fact, this would also be the behavior when we go functional: anything that is not explicitly marked as input or output will be up to the executor for lifetime management. See the comments in net_simple_refcount.cc for details. Reviewed By: dzhulgakov Differential Revision: D12855489 fbshipit-source-id: 594a47a786305d595fd505b6700864dd1d9c72aa
59 lines
2 KiB
C++
59 lines
2 KiB
C++
#ifndef CAFFE2_CORE_NET_SIMPLE_REFCOUNT_H_
|
|
#define CAFFE2_CORE_NET_SIMPLE_REFCOUNT_H_
|
|
|
|
#include <vector>
|
|
|
|
#include "c10/util/Registry.h"
|
|
#include "caffe2/core/common.h"
|
|
#include "caffe2/core/logging.h"
|
|
#include "caffe2/core/net.h"
|
|
#include "caffe2/core/net_simple.h"
|
|
#include "caffe2/core/tensor.h"
|
|
#include "caffe2/core/workspace.h"
|
|
#include "caffe2/proto/caffe2_pb.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
// SimpleRefcountNet is an implementation that adds an additional abstraction
|
|
// on top of SimpleRefCountNet: it tracks all the tensors and for those that are
|
|
// considered internal/temporary, delete them once their refcount go to zero.
|
|
// In the context of a simple static run, this can be carried out during
|
|
// construction time: we will do a pass through the network and track what
|
|
// blobs we need to do reset on, after the execution of every op.
|
|
//
|
|
// To identify which blob is considered temporary, we employ the following
|
|
// strategy: any blob that is
|
|
// (1) consumed but not produced by ops in the net, or
|
|
// (2) produced but not consumed by ops in the net, or
|
|
// (3) is marked as external_output in the protobuf
|
|
// will NOT be considered temporary.
|
|
//
|
|
// In the long run, we should design proper functional interfaces so that
|
|
// nets are less imperative and more functional.
|
|
//
|
|
// Also, for now, SimpleRefCountNet should only be used for benchmarking
|
|
// purposes and not product use, since it is not going to provide better
|
|
// performance gain, and is implicitly incompatible with the contract that
|
|
// earlier Nets expose - that all intermediate blobs are visible to the users.
|
|
class SimpleRefCountNet final : public SimpleNet {
|
|
public:
|
|
SimpleRefCountNet(
|
|
const std::shared_ptr<const NetDef>& net_def,
|
|
Workspace* ws);
|
|
|
|
protected:
|
|
bool Run() override;
|
|
|
|
using SimpleNet::operators_;
|
|
|
|
private:
|
|
// The list of blobs to delete when each operator finishes its run.
|
|
// This will be populated during construction time.
|
|
vector<vector<Blob*>> delete_list_;
|
|
|
|
C10_DISABLE_COPY_AND_ASSIGN(SimpleRefCountNet);
|
|
};
|
|
|
|
} // namespace caffe2
|
|
|
|
#endif // CAFFE2_CORE_NET_SIMPLE_REFCOUNT_H_
|