libtorch demo

这个libtorch1.1可以调用

  

   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

       model = ShuffleNetV2(scale=0.5, in_channels=3, c_tag=0.5, num_classes=2, activation=nn.ReLU,
     
    
   

    

     

    
    

     

       SE=False, residual=False)
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

       model.eval()
     
    
   

    

     

    
    

     

       model=model#.cuda()
     
    
   

    

     

    
    

     

       x = torch.rand(1, 3, 352, 352)#.cuda()
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

       traced_script_module = torch.jit.trace(model, x)
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

       traced_script_module.save("c++_model.pt")
     
    
  
 

这个c++ 能调用：

  

   

    

     

    
    

     

      class ScriptModel(torch.jit.ScriptModule):
     
    
   

    

     

    
    

     
 def __init__(self, N, M):
     
    
   

    

     

    
    

     

       super(ScriptModel, self).__init__()
     
    
   

    

     

    
    

     

       self.weight = torch.nn.Parameter(torch.rand(N, M))
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

       @torch.jit.script_method
     
    
   

    

     

    
    

     
 def forward(self, input):
     
    
   

    

     

    
    

     
 if input.sum() > 0:
     
    
   

    

     

    
    

     

       output = self.weight.mv(input)
     
    
   

    

     

    
    

     
 else:
     
    
   

    

     

    
    

     

       output = self.weight + input
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

       ret
     
    
  
 

文章来源: blog.csdn.net，作者：网奇，版权归原作者所有，如需转载，请联系作者。

原文链接：blog.csdn.net/jacke121/article/details/93385470