mxnet GPU加速计算

  

   

    

     

    
    

     

      #coding:utf-8
     
    
   

    

     

    
    

     

      '''
     
    
   

    

     

    
    

     

      NDArray在GPU上计算
     
    
   

    

     

    
    

     

      '''
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      from mxnet import nd
     
    
   

    

     

    
    

     

      from mxnet.gluon import nn
     
    
   

    

     

    
    

     

      import mxnet as mx
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      a = nd.array([1,2,3],ctx=mx.gpu())
     
    
   

    

     

    
    

     

      b = nd.zeros((3,2),ctx=mx.gpu())
     
    
   

    

     

    
    

     

      x = nd.array([1,2,3])
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      y = x.copyto(mx.gpu())
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      z = x.as_in_context(mx.gpu())
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      print('a = ',a)
     
    
   

    

     

    
    

     

      print('b = ',b)
     
    
   

    

     

    
    

     

      print('x = ',x)
     
    
   

    

     

    
    

     

      print('y = ',y)
     
    
   

    

     

    
    

     

      print('z = ',z)
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      # Gluon的GPU计算
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      net = nn.Sequential()
     
    
   

    

     

    
    

     

      net.add(nn.Dense(1))
     
    
   

    

     

    
    

     

      net.initialize(ctx = mx.gpu())
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      data = nd.random.uniform(shape=(3,2),ctx=mx.gpu())
     
    
   

    

     

    
    

     

      print(net(data))
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      # 模型的参数也储存在GPU上
     
    
   

    

     

    
    

     

      print(net[0].weight.data())
     
    
  
 

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

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