pytorch 30 torchscipt(jit模型)与onnx输入单个int数值与if分支控制

通常来说,pytorch的模型输入和输出都应该是tensor。但是,有的时候,需要传入额外的参数控制模型的执行流程,则需要考虑输入int、bool等类型的变量控制模型的forward的过程。如博主前面写的文章pytorch 20 基于pytorch实现腐蚀膨胀、开运算闭运算等形态学操作_万里鹏程转瞬至的博客-CSDN博客_pytorch 腐蚀  

该文章使用maxpool实现腐蚀膨胀等系列操作。其中的控制maxpool kernel size就需要传入int数值进行控制,以确保按照算法需求动态的进行腐蚀膨胀操作。下面以多输入多输出模型为例,实现模型输入单个int与if分支控制。并用onnxruntime进行了严格的测试,保证了onnx模型也支持输入单个int和if分支控制。

1、创建带if分支的多输入模型

该模型通过传入的flag参数执行两个不同的分支,也急速根据flag值进行逻辑运算得到bool值,然后进行执行分支。执行分支是两个切片操作,如果传入的flag大于0,则将x1_in的第一个值赋给out1的切片区域;如果传入的flag小于0,则将x2_in的第一个值赋给out2的切片区域。

import torch
import torch.nn as nn
import torch.nn.functional as F
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        k_num=128
        sel

猜你喜欢

转载自blog.csdn.net/a486259/article/details/125338478