pytorch topk

可以最大或最小排序：

  

   

    

     

    
    

     

      import  numpy as np
     
    
   

    

     

    
    

     

      import torch
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      import torch
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      output = torch.tensor([[-5.4783, 0.2298],
     
    
   

    

     

    
    

     

       [-4.2573, -0.4794],
     
    
   

    

     

    
    

     

       [-0.1070, -5.1511],
     
    
   

    

     

    
    

     

       [-0.1785, -4.3339]])
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      # x = torch.arange(0, 10).resize_((2, 5))
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      print(output.size())
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      print(output[:,1])
     
    
   

    

     

    
    

     

      topk, indices = torch.topk(output[:,1],2,largest=False)
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      print(topk)
     
    
   

    

     

    
    

     

      print(indices)
     
    
   

    

     

    
    

     

      print(output[indices])
     
    
  
 

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

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