Repository: pytorch
Issue: [caffe2] Memory usage
Labels: caffe2

Hi. I am getting unexpected big amount of memory usage when running onnx models in Caffe2 on C++. 

For example for SSD with VGG feature extractor it gets as high as 1.2Gb, while same model with almost equal size in serialized form (about 104Mb) in TensorFlow consume 130Mb. And i getting similar behavior for other models I tried.

Just wanna know is that expected to memory usage be such high in Caffe2 or I doing something wrong?

Here is my code to evaluate model on Caffe2:
```
#include <string>
#include <vector>
#include <caffe2/onnx/backend.h>
#include <caffe2/core/init.h>
#include <caffe2/utils/proto_utils.h>
#include <onnx/proto_utils.h>
#include <onnx/onnx_ONNX_NAMESPACE.pb.h>
#include <fstream>

static int sleeptime = 1;
static int res = 300;

ONNX_NAMESPACE::ModelProto read_model(std::string path){
    std::ifstream model_file(path, std::ios::binary|std::ios::ate);

    auto size = model_file.tellg();
    model_file.seekg(0, model_file.beg);

    ONNX_NAMESPACE::ModelProto model_proto;
    std::vector<char> model_bin_str(size);

    model_file.readsome(model_bin_str.data(), size);

    ONNX_NAMESPACE::ParseProtoFromBytes(&model_proto, model_bin_str.data(), size);

    model_file.close();

    return model_proto;
}

void f(std::string model_path){

    caffe2::onnx::Caffe2Backend backend;
    auto back = std::unique_ptr<caffe2::onnx::Caffe2BackendRep>(backend.Prepare(read_model(model_path).SerializeAsString(),"CPU :0",{}));

    std::cout << "Backend prepared" << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(sleeptime));

    caffe2::Predictor::TensorVector inputs, outputs;
    caffe2::TensorCPU input_tensor;
    std::vector<caffe2::TIndex> tensor_size = {1, 3, res, res};
    input_tensor.Resize(tensor_size);
    auto tensor_data = input_tensor.mutable_data<float>();
    auto tensor_size_ = std::accumulate(tensor_size.begin(), tensor_size.end(), 1, std::multiplies<int>());
    for (int i = 0; i < tensor_size_; ++i) {
        *(tensor_data+i) = 1.f;
    }
    inputs.push_back(&input_tensor);

    std::cout << "Input prepared" << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(sleeptime));

    back->Run(inputs, &outputs);

    std::cout << "Model inferenced" << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(sleeptime));

    inputs.clear();
    outputs.clear();

    input_tensor.Resize(tensor_size);
    for (int i = 0; i < tensor_size_; ++i) {
        *(tensor_data+i) = 0.f;
    }

    inputs.push_back(&input_tensor);

    std::cout << "Input prepared" << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(sleeptime));

    back->Run(inputs, &outputs);

    std::cout << "Model inferenced" << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(sleeptime));
}

int main(int argc, char *argv[])
{
    sleeptime = std::stoi(std::string(argv[2]));
    res = std::stoi(std::string(argv[3]));
    
    std::cout << "Started" << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(sleeptime));
    
    caffe2::GlobalInit();
    
    f(std::string(argv[1]));
    
    std::cout << "Pass done" << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(sleeptime));
    
    caffe2::ShutdownProtobufLibrary();
    
    return 0;
}
```

And here is my cmake to build it
```
project(caffe2_onnx_mem)
cmake_minimum_required(VERSION 3.5)
set(CMAKE_CXX_STANDARD 11)

include(${CMAKE_ROOT}/Modules/ExternalProject.cmake)

set(CAFFE2_DIR ${CMAKE_CURRENT_BINARY_DIR}/caffe2)
set(CAFFE2_PREFIX ${CMAKE_CURRENT_BINARY_DIR}/Caffe2-prefix)
set(CAFFE2_BUILD_DIR ${CAFFE2_PREFIX}/src/Caffe2-build)
set(CAFFE2_SRC_DIR ${CAFFE2_PREFIX}/src/Caffe2/)

ExternalProject_Add(
        Caffe2
        UPDATE_COMMAND ""
        GIT_REPOSITORY "https://github.com/pytorch/pytorch"
        GIT_TAG "v0.4.1"
        CMAKE_ARGS
                -DONNX_NAMESPACE=ONNX_NAMESPACE
                -DBUILD_BINARY=OFF
                -DUSE_CUDA=OFF
                -DBUILD_PYTHON=OFF
                -DUSE_PROF=OFF
                -DUSE_ATEN=OFF
                -DUSE_OPENCV=OFF
                -DUSE_LMDB=OFF
                -DUSE_LEVELDB=OFF
                -DUSE_GLOO=OFF
                -DUSE_GLOG=OFF
                -DUSE_GFLAGS=OFF
                -DUSE_NNPACK=ON
                -DUSE_MKL=OFF
                -DUSE_MKLML=OFF
                -DUSE_IDEEP=OFF
                -DUSE_NATIVE_ARCH=ON
                -DUSE_LITE_PROTO=OFF
                -DBUILD_SHARED_LIBS=OFF
                -DCMAKE_INSTALL_PREFIX=${CAFFE2_DIR}
        INSTALL_COMMAND make install
        )

add_executable($caffe2_onnx_mem "c2_onnx_mem.cpp")

target_include_directories(caffe2_onnx_mem PUBLIC
        ${CAFFE2_DIR}/include
        CAFFE2_SRC_DIR
        )

target_link_libraries(caffe2_onnx_mem
        -Wl,--whole-archive
        ${CAFFE2_DIR}/lib/libcaffe2.a
        ${CAFFE2_BUILD_DIR}/lib/libcaffe2_protos.a
        -Wl,--no-whole-archive
        ${CAFFE2_DIR}/lib/libcpuinfo.a
        ${CAFFE2_DIR}/lib/libnomnigraph.a
        ${CAFFE2_DIR}/lib/libonnx.a
        ${CAFFE2_DIR}/lib/libnnpack.a
        ${CAFFE2_DIR}/lib/libonnx_proto.a
        ${CAFFE2_BUILD_DIR}/lib/libonnxifi_loader.a
        ${CAFFE2_DIR}/lib/libprotobuf.a
        ${CAFFE2_DIR}/lib/libprotoc.a
        pthread
        dl
        )
```

I measuring memory usage with
```
watch -n 0.1 'pidof ./caffe2_onnx_mem | xargs -i cat /proc/{}/status | grep "RssAnon"'
```

And here is the python code that i using to make a model
```
import torch

class ConvBnAct(torch.nn.Module):
    def __init__(self, inp, out, stride=1, kernel=3, groups=1, dilations=1,bias=False, pad=None):
        super().__init__()
        
        if pad is not None:
            is_pad = pad
        else:
            is_pad = True if kernel != 1 else False
        
        self.conv = torch.nn.Conv2d(inp,out,kernel,stride,int(is_pad),dilations,groups,bias)
        self.bn = torch.nn.BatchNorm2d(out)
        self.relu = torch.nn.LeakyReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x
    
class ConvAct(torch.nn.Module):
    def __init__(self, inp, out, stride=1, kernel=3, groups=1, dilations=1,bias=False, pad=None):
        super().__init__()
        
        if pad is not None:
            is_pad = pad
        else:
            is_pad = True if kernel != 1 else False
        
        self.conv = torch.nn.Conv2d(inp,out,kernel,stride,int(is_pad),dilations,groups,bias)
        self.relu = torch.nn.LeakyReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x
    
class VGGStage(torch.nn.Module):
    def __init__(self, input_depth, output_depth, conv_count=2, pool=True):
        super().__init__()
        
        self.ops = []
        input_depth_ = input_depth
        
        if pool:
            self.ops.append(torch.nn.MaxPool2d(3,stride=2,padding=1))
        
        for i in range(conv_count):
            self.ops.append(ConvBnAct(input_depth_, output_depth))
            input_depth_ = output_depth
                
        self.ops = torch.nn.Sequential(
            *self.ops
        )
        
    def forward(self, x):
        return self.ops(x)

class VGG16_SSD(torch.nn.Module):
    def __init__(self, cls_count=20):
        super().__init__()
        stage_defs = [
            (3,64,2, False),
            (64,128,2, True),
            (128,256,3, True),
            (256,512,3, True),
            (512,512,3, True)
        ]
        
        bbox_counts=[4,6,6,6,4]
        inputs=[512,1024,512,256,256]
        
        self.stages = []
        
        for inp, out, num, pool in stage_defs:
            self.stages.append(VGGStage(inp, out, num, pool))
        
        self.stages += [
            ConvBnAct(512,1024,kernel=3),
            ConvBnAct(1024,1024,kernel=1),
            torch.nn.Sequential(
                ConvBnAct(1024,256,kernel=1),
                ConvBnAct(256,512,kernel=3, stride=2)
            ),
            torch.nn.Sequential(
                ConvBnAct(512,128,kernel=1),
                ConvBnAct(128,256,kernel=3, stride=2)
            ),
            torch.nn.Sequential(
                ConvBnAct(256,128,kernel=1, pad=False),
                ConvBnAct(128,256,kernel=3, pad=False)
            ),
            torch.nn.Sequential(
                ConvBnAct(256,128,kernel=1, pad=False),
                ConvAct(128,256,kernel=3, pad=False)
            ),
        ]
            
        self.stages = torch.nn.ModuleList(self.stages)
        
        self.extra_convs = torch.nn.ModuleList([
           ConvBnAct(input_size,bbox_cnt * (4 + cls_count),kernel=3) for bbox_cnt, input_size in zip(bbox_counts, inputs) 
        ] + [
            ConvAct(256,4,kernel=1)
        ])
    
    
    def forward(self, x):
        
        outputs = []
        
        for stage in self.stages[:4]:
            x = stage(x)
        
        outputs.append(x)
        
        for stage in self.stages[4:7]:
            x = stage(x)
        
        outputs.append(x)
        
        for stage in self.stages[7:]:
            x = stage(x)
            outputs.append(x)
        
        detections = tuple(
            c(o) for o, c in zip(outputs, self.extra_convs)
        )
        
        return detections

net = VGG16_SSD()
tensor = torch.ones((1,3,300,300))
out = net(tensor)
torch.onnx.export(net, tensor, "./vggssd.onnx", verbose=False, input_names=['Input'])        

```