问题描述
有时候，我建立了一个基于function ParentComponent(){
// do sth to get the userinfo here
return(
<ParentComponent>
{userinfo? <Profile userinfo={userinfo} /> : null}
<ParentComponent/>
)
}
的机器学习“图书馆”作为学校的作业。它纯粹基于numpy，但现在我想将其转换为JAX。我在调整反向传播过程时遇到了麻烦。
该库模拟pytorch，因此每一层都是具有numpy和forward方法的类。在backward中，我的线性层是
numpy现在我很难将其转换为JAX。我试图定义一个class Linear:
def __init__(self,in_features: int,out_features: int):
self.W = np.random.randn(in_features,out_features)
self.b = np.random.randn(out_features)
def forward(self,x: Tensor) -> Tensor:
self.input = x
return x @ self.W + self.b
def backward(self,grad: Tensor) -> Tensor:
# in_feat by batch_size @ batch_size by out_feat
self.dydw = self.input.T @ grad
# we sum across batches and get shape (out_features)
self.dydb = grad.sum(axis=0)
# output must be of shape (batch_size,out_features)
return grad @ self.W.T
方法
__matmul__然后使用def __matmul__(
self,weight: Tensor,input_: Tensor,bias: Tensor
) -> Tensor:
return jnp.dot(weight,input_.T) + bias
jax.grad不起作用。如果仅在def __grad__(self,) -> Tuple[Callable,Callable,Callable]:
return (
grad(self.__matmul__,argnums=0),grad(self.__matmul__,argnums=1),argnums=2),)
的输出上使用grad，则结果似乎不正确。在这里使用JAX的autograd的正确方法是什么？
解决方法
暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！
如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。
小编邮箱:dio#foxmail.com (将#修改为@）