• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2021 Huawei Technologies Co., Ltd
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #ifndef MINDSPORE_LITE_TOOLS_OPTIMIZER_FORMAT_TO_NCHW_FORMAT_H_
18 #define MINDSPORE_LITE_TOOLS_OPTIMIZER_FORMAT_TO_NCHW_FORMAT_H_
19 
20 #include "tools/optimizer/format/to_format_base.h"
21 
22 namespace mindspore {
23 namespace opt {
24 class ToNCHWFormat : public ToFormatBase {
25  public:
26   explicit ToNCHWFormat(FmkType fmk_type = converter::kFmkTypeMs, bool train_flag = false,
27                         ModelType save_type = kMindIR)
28       : ToFormatBase(fmk_type, train_flag, save_type, "ToNCHWFormat") {
29     format_ = mindspore::NCHW;
30   }
31   ~ToNCHWFormat() = default;
32 
33   STATUS GetTransNodeFormatType(const CNodePtr &cnode, opt::TransTypePair *trans_info) override;
34 
35  protected:
36   STATUS DecideConvWeightSrcAndDstFormat(const CNodePtr &cnode, schema::Format *src_format,
37                                          schema::Format *dst_format) override;
38 };
39 }  // namespace opt
40 }  // namespace mindspore
41 
42 #endif  // MINDSPORE_LITE_TOOLS_OPTIMIZER_FORMAT_TO_NCHW_FORMAT_H_
43