pytorch Flatten展平

经过测试，cpu模式中，用Module的方式比直接在推理中展平平均快1~5ms 

  

   

    

     

    
    

     

      from torch import nn
     
    
   

    

     

    
    

     

      import torch
     
    
   

    

     

    
    

     

      import math
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      class Flatten(Module):
     
    
   

    

     

    
    

     
 def forward(self, input):
     
    
   

    

     

    
    

     
 return input.view(input.size(0), -1)
     
    
  
 

直接展平： 

  

   

    

     

    
    

     

        out = self.avgpool(out)
     
    
   

    

     

    
    

     
 # out = self.conv_6_flatten(out)
     
    
   

    

     

    
    

     

       out = out.view(out.size(0), -1)
     
    
  
 

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

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