io.h 2.2 KB
Newer Older
W
wangliu 已提交
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
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once

#include <memory.h>
#include <string>
#include <vector>

#include "common/types.h"
#include "framework/lod_tensor.h"
#include "framework/operator.h"
#include "framework/paddle_mobile_object.h"
#include "framework/program/program.h"
#include "framework/tensor.h"

namespace paddle_mobile {

template <typename Dtype, Precision P = Precision::FP32>
class Loader : PaddleMobileObject {
 public:
L
liuruilong 已提交
33 34
  const framework::Program<Dtype, P> Load(const std::string &dirname,
                                          bool optimize = true);
W
wangliu 已提交
35 36 37 38 39 40 41 42 43 44 45 46 47 48

 private:
  void LoadVar(framework::Variable *variable,
               const framework::VarDesc &var_desc,
               const std::string &file_path);
};

template <typename Dtype, Precision P = Precision::FP32>
class Executor {
 public:
  typedef typename PrecisionTrait<P>::ptype Ptype;

  Executor() = default;

L
liuruilong 已提交
49 50
  Executor(const framework::Program<Dtype> p, int batch_size = 1,
           bool use_optimize = true);
W
wangliu 已提交
51

L
liuruilong 已提交
52
  //  std::shared_ptr<framework::Tensor> Predict(framework::Tensor &t);
W
wangliu 已提交
53

L
liuruilong 已提交
54
  std::vector<Ptype> Predict(const std::vector<Ptype> &input,
W
wangliu 已提交
55 56 57 58 59 60 61 62 63
                             const std::vector<int64_t> &dims);

 protected:
  void InitMemory();
  void LoadMemory(const framework::VarDesc var_desc,
                  framework::LoDTensor *tensor, const std::string &file_path);
  framework::Program<Dtype> program_;
  int batch_size_ = 1;
  std::shared_ptr<framework::ProgramDesc> to_predict_program_;
L
liuruilong 已提交
64
  void Predict(const framework::Tensor &t, int block_id);
W
wangliu 已提交
65 66 67 68 69 70 71
  std::map<framework::BlockDesc,
           std::vector<std::shared_ptr<framework::OperatorBase<Dtype>>>>
      ops_of_block_;
  bool use_optimize_ = false;
};

}  // namespace paddle_mobile