dm-haiku 用法

pip install jax jaxlib dm-haiku，2021.02.16只支持linux

  

   

    

     

    
    

     

      class MyDropout(hk.Module):
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     
  def __init__(self, rate=0.5, name=None):
     
    
   

    

     

    
    

     

       super().__init__(name=name)
     
    
   

    

     

    
    

     

       self.rate = rate
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     
  def __call__(self, x):
     
    
   

    

     

    
    

     

       key = hk.next_rng_key()
     
    
   

    

     

    
    

     

       p = jax.random.bernoulli(key, 1.0 - self.rate, shape=x.shape)
     
    
   

    

     

    
    

     
 return x * p / (1.0 - self.rate)
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      forward = hk.transform(lambda x: MyDropout()(x))
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      key1, key2 = jax.random.split(jax.random.PRNGKey(42), 2)
     
    
   

    

     

    
    

     

      params = forward.init(key1, x)
     
    
   

    

     

    
    

     

      prediction = forward.apply(params, key2, x)
     
    
  
 

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

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