1 // Tencent is pleased to support the open source community by making ncnn available. 2 // 3 // Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved. 4 // 5 // Licensed under the BSD 3-Clause License (the "License"); you may not use this file except 6 // in compliance with the License. You may obtain a copy of the License at 7 // 8 // https://opensource.org/licenses/BSD-3-Clause 9 // 10 // Unless required by applicable law or agreed to in writing, software distributed 11 // under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR 12 // CONDITIONS OF ANY KIND, either express or implied. See the License for the 13 // specific language governing permissions and limitations under the License. 14 15 #include "pass_level1.h" 16 17 #include "../utils.h" 18 19 namespace pnnx { 20 21 class BatchNorm1d : public FuseModulePass 22 { 23 public: match_type_str() const24 const char* match_type_str() const 25 { 26 return "__torch__.torch.nn.modules.batchnorm.BatchNorm1d"; 27 } 28 type_str() const29 const char* type_str() const 30 { 31 return "nn.BatchNorm1d"; 32 } 33 write(Operator * op,const std::shared_ptr<torch::jit::Graph> & graph,const torch::jit::Module & mod) const34 void write(Operator* op, const std::shared_ptr<torch::jit::Graph>& graph, const torch::jit::Module& mod) const 35 { 36 const torch::jit::Node* bn = find_node_by_kind(graph, "aten::batch_norm"); 37 38 const auto& running_mean = mod.attr("running_mean").toTensor(); 39 const auto& running_var = mod.attr("running_var").toTensor(); 40 41 op->params["num_features"] = running_mean.size(0); 42 op->params["eps"] = bn->namedInput("eps"); 43 op->params["affine"] = mod.hasattr("weight") && mod.hasattr("bias"); 44 45 op->attrs["running_mean"] = running_mean; 46 op->attrs["running_var"] = running_var; 47 if (mod.hasattr("weight") && mod.hasattr("bias")) 48 { 49 op->attrs["weight"] = mod.attr("weight").toTensor(); 50 op->attrs["bias"] = mod.attr("bias").toTensor(); 51 } 52 } 53 }; 54 55 REGISTER_GLOBAL_PNNX_FUSE_MODULE_PASS(BatchNorm1d) 56 57 } // namespace pnnx 58