Skip to content

Commit

Permalink
support torch.t to ncnn (#4940)
Browse files Browse the repository at this point in the history
  • Loading branch information
XiaBing992 committed Aug 14, 2023
1 parent e80fcbc commit 070a6d4
Show file tree
Hide file tree
Showing 4 changed files with 117 additions and 0 deletions.
1 change: 1 addition & 0 deletions tools/pnnx/src/CMakeLists.txt
Expand Up @@ -537,6 +537,7 @@ set(pnnx_pass_ncnn_SRCS
pass_ncnn/torch_prod.cpp
pass_ncnn/torch_squeeze.cpp
pass_ncnn/torch_sum.cpp
pass_ncnn/torch_t.cpp
pass_ncnn/torch_transpose.cpp
pass_ncnn/torch_unsqueeze.cpp
pass_ncnn/torchvision_DeformConv2d.cpp
Expand Down
54 changes: 54 additions & 0 deletions tools/pnnx/src/pass_ncnn/torch_t.cpp
@@ -0,0 +1,54 @@
// Tencent is pleased to support the open source community by making ncnn available.
//
// Copyright (C) 2023 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// 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.

#include "pass_ncnn.h"

namespace pnnx {

namespace ncnn {

class torch_t : public GraphRewriterPass
{
public:
const char* match_pattern_graph() const
{
return R"PNNXIR(7767517
3 2
pnnx.Input input 0 1 input
torch.t op_0 1 1 input out
pnnx.Output output 1 0 out
)PNNXIR";
}

const char* type_str() const
{
return "Permute";
}

const char* name_str() const
{
return "t";
}

void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
{
op->params["0"] = 1;
}
};

REGISTER_GLOBAL_PNNX_NCNN_GRAPH_REWRITER_PASS(torch_t, 20)

} // namespace ncnn

} // namespace pnnx
1 change: 1 addition & 0 deletions tools/pnnx/tests/ncnn/CMakeLists.txt
Expand Up @@ -153,6 +153,7 @@ pnnx_ncnn_add_test(torch_norm)
pnnx_ncnn_add_test(torch_permute)
pnnx_ncnn_add_test(torch_prod)
pnnx_ncnn_add_test(torch_sum)
pnnx_ncnn_add_test(torch_t)
pnnx_ncnn_add_test(torch_squeeze)
pnnx_ncnn_add_test(torch_stack)
pnnx_ncnn_add_test(torch_tensor_split)
Expand Down
61 changes: 61 additions & 0 deletions tools/pnnx/tests/ncnn/test_torch_t.py
@@ -0,0 +1,61 @@
# Tencent is pleased to support the open source community by making ncnn available.
#
# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
#
# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# 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.

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()

def forward(self, x, y):
x = torch.t(x)
y = torch.t(y)
x = F.relu(x)
y = F.relu(y)
return x, y

def test():
net = Model()
net.eval()

torch.manual_seed(0)
x = torch.rand(3)
y = torch.rand(5, 9)

a = net(x, y)

# export torchscript
mod = torch.jit.trace(net, (x, y))
mod.save("test_torch_t.pt")

# torchscript to pnnx
import os
os.system("../../src/pnnx test_torch_t.pt inputshape=[3],[5,9]")

# ncnn inference
import test_torch_t_ncnn
b = test_torch_t_ncnn.test_inference()

for a0, b0 in zip(a, b):
if not torch.allclose(a0, b0, 1e-4, 1e-4):
return False
return True

if __name__ == "__main__":
if test():
exit(0)
else:
exit(1)

0 comments on commit 070a6d4

Please sign in to comment.