reid from https://github.com/michuanhaohao/reid-strong-baseline
zhangmeng
2020-01-17 f7c4a3cfd07adede3308f8d9d3d7315427d90a7c
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
#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_