{"id":10008044,"date":"2023-07-05T23:59:04","date_gmt":"2023-07-05T15:59:04","guid":{"rendered":"https:\/\/ityet.com\/?p=10008044"},"modified":"2023-07-05T23:59:04","modified_gmt":"2023-07-05T15:59:04","slug":"%e6%9c%ba%e5%99%a8%e5%ad%a6%e4%b9%a0%e6%b4%9e%e5%af%9f-jax%ef%bc%8c%e6%9c%ba%e5%99%a8%e5%ad%a6%e4%b9%a0%e9%a2%86%e5%9f%9f%e7%9a%84%e6%96%b0%e9%9d%a2%e5%ad%94","status":"publish","type":"post","link":"https:\/\/itmax.cn\/?p=10008044","title":{"rendered":"\u673a\u5668\u5b66\u4e60\u6d1e\u5bdf | JAX\uff0c\u673a\u5668\u5b66\u4e60\u9886\u57df\u7684\u201c\u65b0\u9762\u5b54\u201d"},"content":{"rendered":"<p>\u5728\u4e4b\u524d\u7684\u300a\u673a\u5668\u5b66\u4e60\u6d1e\u5bdf\u300b\u7cfb\u5217\u6587\u7ae0\u4e2d\uff0c\u6211\u4eec\u5206\u522b\u9488\u5bf9\u4e8e\u591a\u6a21\u6001\u673a\u5668\u5b66\u4e60\u548c\u5206\u5e03\u5f0f\u8bad\u7ec3\u3001\u65e0\u670d\u52a1\u5668\u63a8\u7406\u8fdb\u884c\u4e86\u89e3\u8bfb\uff0c\u672c\u6587\u5c06\u4e3a\u60a8\u91cd\u70b9\u4ecb\u7ecd JAX \u7684\u53d1\u5c55\u5e76\u5256\u6790\u5176\u6f14\u53d8\u548c\u52a8\u673a\u3002\u4e0b\u9762\uff0c\u5c31\u8ba9\u6211\u4eec\u6765\u8ba4\u8bc6\u4e00\u4e0b JAX \u8fd9\u4e00\u65b0\u5d1b\u8d77\u7684\u6df1\u5ea6\u5b66\u4e60\u6846\u67b6\u2014\u2014<\/p>\n<table>\n<tbody>\n<tr>\n<td><a href=\"https:\/\/link.segmentfault.com\/?enc=L8TZgt62ysbTEvsqhWg6Lw%3D%3D.tGm3OerNXjCf7RLp%2BnBkL6UvYhpqMpPVgp9ZP0MRumM30Nzjuoguqm3Bvzxy1dLMIvyx0OhEGCOClNYBlGrLUdypLQ0t0fz5GWx2xMD72JrtTKldsVRJkP1okONT%2BmUh\" rel=\"nofollow\">\u4e9a\u9a6c\u900a\u4e91\u79d1\u6280\u5f00\u53d1\u8005\u793e\u533a<\/a>\u4e3a\u5f00\u53d1\u8005\u4eec\u63d0\u4f9b\u5168\u7403\u7684\u5f00\u53d1\u6280\u672f\u8d44\u6e90\u3002\u8fd9\u91cc\u6709\u6280\u672f\u6587\u6863\u3001\u5f00\u53d1\u6848\u4f8b\u3001\u6280\u672f\u4e13\u680f\u3001\u57f9\u8bad\u89c6\u9891\u3001\u6d3b\u52a8\u4e0e\u7ade\u8d5b\u7b49\u3002\u5e2e\u52a9\u4e2d\u56fd\u5f00\u53d1\u8005\u5bf9\u63a5\u4e16\u754c\u6700\u524d\u6cbf\u6280\u672f\uff0c\u89c2\u70b9\uff0c\u548c\u9879\u76ee\uff0c\u5e76\u5c06\u4e2d\u56fd\u4f18\u79c0\u5f00\u53d1\u8005\u6216\u6280\u672f\u63a8\u8350\u7ed9\u5168\u7403\u4e91\u793e\u533a\u3002\u5982\u679c\u4f60\u8fd8\u6ca1\u6709\u5173\u6ce8\/\u6536\u85cf\uff0c\u770b\u5230\u8fd9\u91cc\u8bf7\u4e00\u5b9a\u4e0d\u8981\u5306\u5306\u5212\u8fc7\uff0c<a href=\"https:\/\/link.segmentfault.com\/?enc=AJUOoBz64tfjbOpkAeLldA%3D%3D.GFwnaQg6jnMSN72pNfG2MF%2FibAmbsRM10is5%2BudIv1v%2Fct4vrDA0C84l7rADXB3z%2F6jnR%2FC4MHuE7nvvFzQm9duHwncIprm4w8AFlMA5D7k2eUK1yTADt36U7WUrHbU98aAgf%2FDt%2FMxcFA1oNTQRMydjJAtfC%2BOPwgtyPghpxlY%3D\" rel=\"nofollow\">\u70b9\u8fd9\u91cc<\/a>\u8ba9\u5b83\u6210\u4e3a\u4f60\u7684\u6280\u672f\u5b9d\u5e93\uff01<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<h3><strong>\u5f00\u6e90\u673a\u5668\u5b66\u4e60\u6846\u67b6\u7684\u6f14\u8fdb<\/strong><\/h3>\n<p><img decoding=\"async\" referrerpolicy=\"no-referrer\" src=\"https:\/\/segmentfault.com\/img\/bVc8FZB\" alt=\"\" title=\"\" loading=\"lazy\"><\/p>\n<p>\u4ece\u8fd9\u5f20 GitHub Star \u8d8b\u52bf\u56fe\u53ef\u4ee5\u770b\u5230\uff0c\u81ea 2019 \u5e74 JAX \u51fa\u73b0\u5230\u5982\u4eca\u4fdd\u6301\u7740\u4e00\u4e2a\u5411\u4e0a\u7684\u629b\u7269\u7ebf\u8d70\u52bf\u3002<\/p>\n<p>\u5728\u8003\u5bdf\u4e00\u4e2a\u5f00\u6e90\u673a\u5668\u5b66\u4e60\u6846\u67b6\u65f6\uff0c\u4f8b\u5982\u5f00\u53d1\u8005\u719f\u77e5\u7684 PyTorch, TensorFlow, MXNet \u7b49\uff0c\u5f80\u5f80\u4f1a\u4ece\u652f\u6301\u6a21\u578b\u7684\u5e7f\u6cdb\u6027\u3001\u90e8\u7f72\u7684\u6210\u719f\u6027\u3001\u751f\u6001\u7cfb\u7edf\u7684\u4e30\u5bcc\u6027\u6765\u5bf9\u5b83\u505a\u4e00\u4e2a\u8bc4\u4f30\uff1a\u5305\u62ec\u662f\u5426\u652f\u6301 Hugging Face \u7b49\u4e3b\u6d41\u6a21\u578b\uff0c\u4ee5\u53ca\u5176\u6846\u67b6\u76f8\u5173\u7814\u7a76\u8bba\u6587\u7684\u6570\u91cf\uff0c\u8fd8\u6709\u5b83\u53ef\u63d0\u4f9b\u590d\u73b0\u4ee3\u7801\u7684\u8bba\u6587\u6570\u91cf\u7b49\u7b49\u3002<\/p>\n<h3><strong>JAX \u7684\u6e90\u8d77<\/strong><\/h3>\n<ul>\n<li>\u4e3a\u4ec0\u4e48 Eager \u6a21\u5f0f\u662f\u5728 TensorFlow 1.4 \u7248\u672c\u4e4b\u540e\u5f15\u5165\u7684\uff1f<\/li>\n<li>Eager \u6a21\u5f0f\u5728 TensorFlow 2.0 \u4e4b\u540e\u53d8\u6210\u4e86\u4e00\u4e2a\u9ed8\u8ba4\u7684\u6267\u884c\u6a21\u5f0f\uff0c\u548c\u539f\u6709\u7684 Graph \u6a21\u5f0f\u7684\u533a\u522b\u662f\u4ec0\u4e48\uff1f<\/li>\n<li>\u2026<\/li>\n<\/ul>\n<p>\u56de\u5f52\u5e76\u7406\u6e05\u8fd9\u4e9b\u5386\u53f2\u95ee\u9898\u6709\u52a9\u4e8e\u5f00\u53d1\u8005\u4e86\u89e3\u673a\u5668\u5b66\u4e60\u7684\u6f14\u53d8\u903b\u8f91\uff0c\u5e76\u4e86\u89e3 JAX \u662f\u5982\u4f55\u5438\u53d6\u4e4b\u524d\u7684\u6559\u8bad\uff0c\u5e2e\u52a9\u5f00\u53d1\u8005\u66f4\u65b9\u4fbf\u5730\u5b9e\u8df5\u6df1\u5ea6\u5b66\u4e60\u6216\u673a\u5668\u5b66\u4e60\u5e94\u7528\u3002<\/p>\n<h3><strong>Eager \u6a21\u5f0f V.S. Graph \u6a21\u5f0f<\/strong><\/h3>\n<p>\u5728 TF \u5f15\u8fdb\u4e86 Eager \u6a21\u5f0f\u4e4b\u540e\uff0c\u5b83\u4f1a\u91c7\u7528\u66f4\u76f4\u89c2\u7684\u754c\u9762\uff0c\u4f7f\u7528\u81ea\u7136\u7684 Python \u4ee3\u7801\u548c\u6570\u636e\u7ed3\u6784\uff0c\u800c\u4e14\u4eab\u53d7\u66f4\u52a0\u4fbf\u643a\u7684\u8c03\u8bd5\uff0c\u5728 Eager \u6a21\u5f0f\u4e2d\u53ef\u4ee5\u901a\u8fc7\u76f4\u63a5\u8c03\u7528\u64cd\u4f5c\u6765\u68c0\u67e5\u548c\u6d4b\u8bd5\u6a21\u578b\uff0c\u800c\u4e4b\u524d Graph \u8fd9\u79cd\u6a21\u5f0f\u6709\u70b9\u7c7b\u4f3c\u4e8e C \u548c C++\uff0c\u5b83\u7684\u7f16\u7a0b\u662f\u5199\u597d\u7a0b\u5e8f\u4e4b\u540e\u8981\u5148\u8fdb\u884c\u7f16\u8bd1\u624d\u80fd\u8fd0\u884c\u3002<\/p>\n<p>Eager \u6a21\u5f0f\u6709\u81ea\u7136\u63a7\u5236\u7684\u6d41\u7a0b\uff0c\u4f7f\u7528 Python \u800c\u4e0d\u662f\u56fe\u63a7\u5236\u6d41\uff0c\u4ee5\u53ca\u652f\u6301 GPU \u548c TPU \u7684\u52a0\u901f\u3002\u505a\u4e3a\u5f00\u53d1\u8005\uff0c\u6211\u4eec\u5e0c\u671b\u53ef\u4ee5\u5ba2\u89c2\u5730\u770b\u5f85\u4e0d\u540c\u7684\u6846\u67b6\uff0c\u800c\u4e0d\u662f\u6bd4\u8f83\u4ed6\u4eec\u7684\u4f18\u52a3\u3002\u503c\u5f97\u601d\u8003\u7684\u4e00\u4e2a\u95ee\u9898\u662f\uff1a\u901a\u8fc7\u4e86\u89e3 TF \u7684 Eager \u6a21\u5f0f\u5bf9\u4e8e Graph \u6a21\u5f0f\u7684\u6539\u8fdb\uff0c\u5b83\u7684\u6539\u8fdb\u903b\u8f91\u548c\u601d\u8def\u5728 JAX \u4e2d\u90fd\u6709\u8eab\u5f71\u3002<\/p>\n<h3><strong>\u4ec0\u4e48\u662f JAX<\/strong><\/h3>\n<p><img decoding=\"async\" referrerpolicy=\"no-referrer\" src=\"https:\/\/segmentfault.com\/img\/bVc8FZF\" alt=\"\" title=\"\" loading=\"lazy\"><\/p>\n<p>JAX \u4f5c\u4e3a\u73b0\u5728\u8d8a\u6765\u8d8a\u6d41\u884c\u7684\u5e93\uff0c\u662f\u4e00\u79cd\u7c7b\u4f3c\u4e8e NumPy\uff08\u4f7f\u7528 Python \u5f00\u6e90\u7684\u6570\u503c\u8ba1\u7b97\u6269\u5c55\u5e93\uff09\u7684\u8f7b\u91cf\u7ea7\u7528\u4e8e\u9635\u5217\u7684\u8ba1\u7b97\u3002JAX \u6700\u5f00\u59cb\u7684\u8bbe\u8ba1\u4e0d\u4ec5\u4ec5\u662f\u4e3a\u4e86\u6df1\u5ea6\u5b66\u4e60\u800c\u8bbe\u8ba1\u7684\uff0c\u6df1\u5ea6\u5b66\u4e60\u53ea\u662f\u5b83\u7684\u4e00\u5c0f\u90e8\u5206\uff0c\u5b83\u63d0\u4f9b\u4e86\u7f16\u5199 NumPy \u7a0b\u5e8f\u7684\u80fd\u529b\uff0c\u8fd9\u4e9b\u7a0b\u5e8f\u53ef\u4ee5\u4f7f\u7528 GPU\/TPU \u81ea\u52a8\u62c6\u5206\u548c\u52a0\u901f\u3002<\/p>\n<p>JAX \u7528\u4e8e\u57fa\u4e8e\u9635\u5217\u7684\u8ba1\u7b97\u65f6\uff0c\u5f00\u53d1\u8005\u65e0\u9700\u4fee\u6539\u4ee3\u7801\u5c31\u53ef\u4ee5\u5728 CPU\/GPU\/ASIC \u4e0a\u540c\u65f6\u8fd0\u884c\uff0c\u5e76\u652f\u6301\u539f\u751f Python \u548c NumPy \u51fd\u6570\u7684\u56db\u79cd\u53ef\u7ec4\u5408\u51fd\u6570\u8f6c\u6362\uff1a<\/p>\n<ul>\n<li>\u81ea\u52a8\u5fae\u5206 (Autodiff)<\/li>\n<li>\u5373\u65f6\u7f16\u8bd1 (JIT compilation)<\/li>\n<li>\u81ea\u52a8\u5411\u91cf\u5316 (Vectorization)<\/li>\n<li>\u4ee3\u7801\u5e76\u884c\u5316 (Parallelization)<\/li>\n<\/ul>\n<h3><strong>JAX \u521d\u4f53\u9a8c<\/strong><\/h3>\n<p>\u6211\u4eec\u53ef\u4ee5\u901a\u8fc7\u4e0b\u9762\u8fd9\u4e2a\u7b80\u5355\u7684\u6d4b\u8bd5\u5bf9\u6bd4 JAX \u548c NumPy \u7684\u8ba1\u7b97\u6027\u80fd\u3002<\/p>\n<p>\u8f93\u5165\u4e00\u4e2a 100 X 100 \u7684\u4e8c\u7ef4\u6570\u7ec4 X\uff0c\u9009\u53d6 ml.g4dn.12xlarge \u8ba1\u7b97\u5b9e\u4f8b\u901a\u8fc7 NumPy \u548c JAX \u5206\u522b\u5bf9\u77e9\u9635\u7684\u524d\u4e09\u6b21\u5e42\u6c42\u548c\uff1a<\/p>\n<pre><code>def fn(x):\n  return x + x*x + x*x*x\nx = np.random.randn(10000, 10000).astype(dtype='float32')\n%timeit -n5 fn(x)\n\n436 ms \u00b1 206 \u00b5s per loop (mean \u00b1 std. dev. of 7 runs, 5 loops each)<\/code><\/pre>\n<p>\u6211\u4eec\u53d1\u73b0\u6b64\u8ba1\u7b97\u5927\u7ea6\u9700\u8981 436 \u6beb\u79d2\u3002\u63a5\u4e0b\u6765\uff0c\u6211\u4eec\u4f7f\u7528 JAX \u5b9e\u73b0\u4ee5\u4e0b\u8ba1\u7b97\uff1a<\/p>\n<pre><code>jax_fn = jit(fn)\nx = jnp.array(x)\n%timeit jax_fn(x).block_until_ready()\n3.67 ms \u00b1 10.7 \u00b5s per loop (mean \u00b1 std. dev. of 7 runs, 1 loop each)\n<\/code><\/pre>\n<p>JAX \u4ec5\u5728 3.67 \u6beb\u79d2\u5185\u6267\u884c\u6b64\u8ba1\u7b97\uff0c\u6bd4 NumPy \u5feb 118 \u500d\u4ee5\u4e0a\u3002\u53ef\u89c1\uff0cJAX \u6709\u53ef\u80fd\u6bd4 NumPy \u5feb\u51e0\u4e2a\u6570\u91cf\u7ea7\uff08\u6ce8\u610f\uff0cJAX \u4f7f\u7528 TPU \u800c NumPy \u6b63\u5728\u4f7f\u7528 CPU\uff09\u3002<\/p>\n<p>*\u4ee5\u4e0a\u4e3a\u4e2a\u4eba\u6d4b\u8bd5\u7ed3\u679c\uff0c\u975e\u5b98\u65b9\u63d0\u4f9b\u7684\u6570\u636e\uff0c\u4ec5\u4f9b\u7814\u7a76\u53c2\u8003<\/p>\n<p>\u5bf9\u6bd4\u6d4b\u8bd5\u7ed3\u679c\u53ef\u5f97\uff0cNumPy \u5b8c\u6210\u8ba1\u7b97\u9700\u8981 436 \u6beb\u79d2\uff0c\u800c JAX \u4ec5\u9700\u8981 3.67 \u6beb\u79d2\uff0c\u8ba1\u7b97\u901f\u5ea6\u76f8\u5dee 100 \u591a\u500d\u3002\u8fd9\u4e2a\u6d4b\u8bd5\u4e5f\u8bf4\u660e\u4e86\u4e3a\u4ec0\u4e48\u5f88\u591a\u5f00\u53d1\u8005\u5bf9\u5b83\u7684\u6027\u80fd\u8d5e\u4e0d\u7edd\u53e3\u3002<\/p>\n<h3><strong>JAX \u7684\u52a8\u673a\u5256\u6790<\/strong><\/h3>\n<p>\u6211\u4eec\u5e0c\u671b\u901a\u8fc7\u56de\u7b54\u8fd9\u4e2a\u95ee\u9898\u6765\u89e3\u8bfb JAX \u7684\u52a8\u673a\uff1a<\/p>\n<p>\u5982\u4f55\u4f7f\u7528 Python \u4ece\u5934\u5f00\u59cb\u5b9e\u73b0\u9ad8\u6027\u80fd\u548c\u53ef\u6269\u5c55\u7684\u6df1\u5ea6\u795e\u7ecf\u7f51\u7edc\uff1f<\/p>\n<h3><strong>\u5728 NumPy \u4e2d\u521b\u5efa\u6df1\u5ea6\u5b66\u4e60\u7cfb\u7edf<\/strong><\/h3>\n<p>\u901a\u5e38\uff0cPython \u7a0b\u5e8f\u5458\u4f1a\u4ece NumPy \u4e4b\u7c7b\u7684\u4e1c\u897f\u5f00\u59cb\uff0c\u56e0\u4e3a\u5b83\u662f\u4e00\u79cd\u719f\u6089\u7684\u3001\u57fa\u4e8e\u6570\u7ec4\u7684\u6570\u636e\u5904\u7406\u8bed\u8a00\uff0c\u5728 Python \u793e\u533a\u4e2d\u5df2\u7ecf\u4f7f\u7528\u4e86\u51e0\u5341\u5e74\u3002\u5982\u679c\u4f60\u60f3\u5728 NumPy \u4e2d\u521b\u5efa\u6df1\u5ea6\u5b66\u4e60\u7cfb\u7edf\uff0c\u4f60\u53ef\u4ee5\u4ece\u9884\u6d4b\u65b9\u6cd5\u5f00\u59cb\u3002<\/p>\n<p>\u8fd9\u91cc\u53ef\u4ee5\u7528\u4e00\u4e2a\u8be6\u7ec6\u7684\u4f8b\u5b50\u8bf4\u660e\u95ee\u9898\uff0c\u4ece NumPy \u4e0a\u7684\u6df1\u5ea6\u5b66\u4e60\u7684\u573a\u666f\u8bf4\u8d77\uff1a<\/p>\n<p><img decoding=\"async\" referrerpolicy=\"no-referrer\" src=\"https:\/\/segmentfault.com\/img\/bVc8F0M\" alt=\"\" title=\"\" loading=\"lazy\"><\/p>\n<p>\u4e0a\u8ff0\u4ee3\u7801\u5c55\u793a\u4e86\u8ba2\u9605\u4e00\u4e2a\u524d\u9988\u7684\u795e\u7ecf\u7f51\u7edc\uff0c\u5b83\u6267\u884c\u4e86\u4e00\u7cfb\u5217\u7684\u70b9\u79ef\u548c\u6fc0\u6d3b\u51fd\u6570\uff0c\u7136\u540e\u5c06\u8f93\u5165\u8f6c\u5316\u4e3a\u67d0\u79cd\u53ef\u4ee5\u5b66\u4e60\u7684\u8f93\u51fa\u3002\u4e00\u65e6\u5b9a\u4e49\u4e86\u8fd9\u6837\u7684\u4e00\u4e2a\u6a21\u578b\uff0c\u63a5\u4e0b\u6765\u9700\u8981\u505a\u5c31\u662f\u8981\u5b9a\u4e49\u635f\u5931\u51fd\u6570\uff0c\u8fd9\u4e2a\u51fd\u6570\u5c06\u4e3a\u4f60\u63d0\u4f9b\u6b63\u5728\u5c1d\u8bd5\u4f18\u5316\u7684\u90a3\u4e9b\u6307\u6807\uff0c\u6765\u9002\u5e94\u6700\u4f73\u7684\u673a\u5668\u5b66\u4e60\u6a21\u578b\u3002\u4f8b\u5982\u4ee5\u4e0a\u4ee3\u7801\u7684\u635f\u5931\u51fd\u6570\u662f\u4ee5\u5747\u65b9\u8bef\u5dee\u635f\u5931\u51fd\u6570 MSE \u4e3a\u4f8b\u3002<\/p>\n<p>\u73b0\u5728\u6211\u4eec\u6765\u5206\u6790\u4e0b\uff1a\u5728\u6df1\u5ea6\u5b66\u4e60\u573a\u666f\u4f7f\u7528 NumPy \u8fd8\u7f3a\u5c11\u4ec0\u4e48\uff1f<\/p>\n<p>\u786c\u4ef6\u52a0\u901f (GPU\/TPU)<\/p>\n<p>\u81ea\u52a8\u5fae\u5206 (autodiff) \u5feb\u901f\u4f18\u5316<\/p>\n<p>\u6dfb\u52a0\u7f16\u8bd1 (Compilation) \u878d\u5408\u64cd\u4f5c<\/p>\n<p>\u5411\u91cf\u5316\u64cd\u4f5c\u6279\u5904\u7406 (batching)<\/p>\n<p>\u5927\u578b\u6570\u636e\u96c6\u5e76\u884c\u5316 (Parallelization)<\/p>\n<p>1\uff09\u786c\u4ef6\u52a0\u901f (GPU\/TPU)\uff1a\u9996\u5148\u6df1\u5ea6\u5b66\u4e60\u9700\u8981\u5927\u91cf\u7684\u8ba1\u7b97\uff0c\u6211\u4eec\u60f3\u5728\u52a0\u901f\u7684\u786c\u4ef6\u4e0a\u8fd0\u884c\u5b83\u3002\u6240\u4ee5\u6211\u4eec\u60f3\u5728 GPU \u548c TPU\/ASIC \u4e0a\u8fd0\u884c\u8fd9\u4e2a\u6a21\u578b\uff0c\u8fd9\u5bf9\u4e8e\u7ecf\u5178\u7684 NumPy \u6765\u8bf4\u6709\u70b9\u56f0\u96be\uff1b<\/p>\n<p>2\uff09\u81ea\u52a8\u5fae\u5206 (autodiff) \u5feb\u901f\u4f18\u5316\uff1a\u63a5\u4e0b\u6765\u6211\u4eec\u60f3\u8981\u505a\u81ea\u52a8\u5fae\u5206\uff0c\u8fd9\u6837\u5c31\u53ef\u4ee5\u6709\u6548\u5730\u62df\u5408\u8fd9\u4e2a\u635f\u5931\u51fd\u6570\uff0c\u800c\u4e0d\u5fc5\u81ea\u5df1\u6765\u5b9e\u73b0\u6570\u503c\u5fae\u5206\uff1b<\/p>\n<p>3\uff09\u7136\u540e\u6211\u4eec\u9700\u8981\u6dfb\u52a0\u7f16\u8bd1 (Compilation)\uff1a\u8fd9\u6837\u4f60\u5c31\u53ef\u4ee5\u5c06\u8fd9\u4e9b\u64cd\u4f5c\u878d\u5408\u5728\u4e00\u8d77\uff0c\u4f7f\u5b83\u4eec\u66f4\u52a0\u9ad8\u6548\uff1b<\/p>\n<p>4\uff09\u5411\u91cf\u5316\u64cd\u4f5c\u6279\u5904\u7406 (Batching)\uff1a\u53e6\u5916\uff0c\u5f53\u6211\u4eec\u7f16\u5199\u4e86\u67d0\u4e9b\u51fd\u6570\u540e\uff0c\u53ef\u80fd\u5e0c\u671b\u5c06\u5176\u5e94\u7528\u4e8e\u591a\u4e2a\u6570\u636e\u7247\u6bb5\uff0c\u800c\u4e0d\u518d\u9700\u8981\u91cd\u5199\u9884\u6d4b\u548c\u635f\u5931\u51fd\u6570\u6765\u5904\u7406\u8fd9\u4e9b\u6279\u91cf\u6570\u636e\uff1b<\/p>\n<p>5\uff09\u5927\u578b\u6570\u636e\u96c6\u5e76\u884c\u5316 (Parallelization)\uff1a\u6700\u540e\uff0c\u5982\u679c\u6211\u4eec\u6b63\u5728\u5904\u7406\u5927\u578b\u6570\u636e\u96c6\uff0c\u4f1a\u5e0c\u671b\u80fd\u591f\u652f\u6301\u8de8\u591a\u4e2a cores \u6216\u591a\u53f0 machines \u505a\u5e76\u884c\u5316\u64cd\u4f5c\u3002<\/p>\n<h3><strong>JAX \u7684\u52a8\u673a\u5256\u6790\uff1aXLA \u548c\u81ea\u52a8\u5b9a\u4f4d<\/strong><\/h3>\n<p>JAX \u975e\u5e38\u91cd\u8981\u7684\u4e00\u4e2a\u52a8\u673a\u5c31\u662f XLA \u548c\u81ea\u52a8\u5b9a\u4f4d\u3002\u8ba9\u6211\u4eec\u6765\u770b\u770b JAX \u53ef\u4ee5\u505a\u4e9b\u4ec0\u4e48\uff0c\u6765\u586b\u8865\u524d\u9762\u5206\u6790\u7684\u5728\u6df1\u5ea6\u5b66\u4e60\u573a\u666f\u4f7f\u7528 NumPy \u8fd8\u7f3a\u5c11\u7684\u529f\u80fd\u3002<br \/><img decoding=\"async\" referrerpolicy=\"no-referrer\" src=\"https:\/\/segmentfault.com\/img\/bVc8F0O\" alt=\"\" title=\"\" loading=\"lazy\"><\/p>\n<p>\u9996\u5148\uff0c\u7528 jax.numpy \u66ff\u6362 numpy \u5bfc\u5165\u6a21\u5757\u3002\u5728\u8bb8\u591a\u60c5\u51b5\u4e0b\uff0cjax.numpy \u4e0e\u7ecf\u5178\u7684 NumPy \u5177\u6709\u76f8\u540c\u7684 API\uff0c\u4f46 jax.numpy \u53ef\u4ee5\u5b8c\u6210\u524d\u9762\u5206\u6790\u65f6\u53d1\u73b0 NumPy \u7f3a\u5c11\uff0c\u4f46\u662f\u5728\u6df1\u5ea6\u5b66\u4e60\u573a\u666f\u5374\u975e\u5e38\u9700\u8981\u7684\u7684\u4e1c\u897f\u3002<\/p>\n<p><strong>JAX \u53ef\u4ee5\u901a\u8fc7 XLA \u540e\u7aef\uff0c\u6765\u81ea\u52a8\u5b9a\u4f4d CPU\u3001GPU \u548c TPU \u6216\u8005 ASIC\uff0c\u4ee5\u4fbf\u5feb\u901f\u8ba1\u7b97\u6a21\u578b\u548c\u7b97\u6cd5\u3002<\/strong><\/p>\n<h3><strong>JAX \u52a8\u673a\u5256\u6790\uff1aAutograd<\/strong><\/h3>\n<p>\u7b2c\u4e8c\u4e2a\u91cd\u8981\u52a8\u673a\u662f Autograd\u3002\u5f00\u53d1\u8005\u53ef\u4ee5\u901a\u8fc7\u4e0b\u9762\u7684\u4ee3\u7801\u8c03\u7528 Autograd \u7248\u672c\uff1a<br \/><img decoding=\"async\" referrerpolicy=\"no-referrer\" src=\"https:\/\/segmentfault.com\/img\/bVc8F0S\" alt=\"\" title=\"\" loading=\"lazy\"><br \/>\u901a\u8fc7 from jax import grad \u6a21\u5757\uff0c\u4f7f\u7528 Autograd \u7684\u66f4\u65b0\u7248\u672c\uff0c<strong>JAX \u53ef\u4ee5\u81ea\u52a8\u5fae\u5206\u539f\u751f Python \u548c NumPy \u51fd\u6570\u3002<\/strong>\u5b83\u53ef\u4ee5\u5904\u7406 Python \u529f\u80fd\u7684\u5927\u5b50\u96c6\uff0c\u5305\u62ec\u5faa\u73af\u3001Ifs\u3001\u9012\u5f52\u7b49\uff0c\u751a\u81f3\u53ef\u4ee5\u63a5\u53d7\u5bfc\u6570\u7684\u5bfc\u6570\u3002<\/p>\n<p>JAX \u63d0\u4f9b\u4e86\u4e00\u7ec4\u53ef\u7ec4\u5408\u7684\u53d8\u6362\uff0c\u5176\u4e2d\u4e4b\u4e00\u662f <strong>grad \u53d8\u6362<\/strong>\u3002<\/p>\n<p>\u4f8b\u5b50\u4e2d\uff0c\u50cf mse_loss \u8fd9\u6837\u7684\u635f\u5931\u51fd\u6570\uff0c\u901a\u8fc7 grad (mse_loss) \u5c06\u5176\u8f6c\u6362\u4e3a\u8ba1\u7b97\u68af\u5ea6\u7684 Python \u51fd\u6570\u3002<\/p>\n<p><strong>Autograd \u7684\u4e3b\u8981\u9884\u671f\u5e94\u7528\u662f\u57fa\u4e8e\u68af\u5ea6\u7684\u4f18\u5316\u3002<\/strong><\/p>\n<p>\u6709\u5173\u66f4\u591a\u4fe1\u606f\uff0c\u8bf7\u67e5\u770b JAX \u6559\u7a0b\u548c\u793a\u4f8b\uff1aHttps:\/\/github.com\/hips\/autograd<\/p>\n<h3><strong>JAX \u52a8\u673a\u5256\u6790\uff1avmap<\/strong><\/h3>\n<p>\u5728\u4f7f\u7528\u68af\u5ea6\u51fd\u6570\u65f6\uff0c\u5f00\u53d1\u8005\u5e0c\u671b\u5c06\u5176\u5e94\u7528\u4e8e\u591a\u4e2a\u6570\u636e\u7247\u6bb5\uff0c\u800c\u5728 JAX \u4e2d\uff0c\u4f60\u4e0d\u518d\u9700\u8981\u91cd\u5199\u9884\u6d4b\u548c\u635f\u5931\u51fd\u6570\u6765\u5904\u7406\u8fd9\u4e9b\u6279\u91cf\u6570\u636e\u3002<\/p>\n<p><img decoding=\"async\" referrerpolicy=\"no-referrer\" src=\"https:\/\/segmentfault.com\/img\/bVc8F0X\" alt=\"\" title=\"\" loading=\"lazy\"><\/p>\n<p>\u5982\u56fe\u4e2d\u4ee3\u7801\u6700\u540e\u4e00\u884c (perexample_grads \u2026) \u6240\u8be0\u91ca\u7684\u90a3\u6837\uff0c\u5982\u679c\u4f60\u901a\u8fc7 vmap transform \u4f20\u9012\u5b83\uff0c\u8fd9\u4f1a\u81ea\u52a8\u5411\u91cf\u5316\u8fd9\u4e2a\u4ee3\u7801\uff0c\u8fd9\u6837\u5c31\u53ef\u4ee5\u5728\u591a\u4e2a\u6279\u6b21\u4e2d\u4f7f\u7528\u76f8\u540c\u7684\u4ee3\u7801\u3002<\/p>\n<h3><strong>JAX \u52a8\u673a\u5256\u6790\uff1ajit<\/strong><\/h3>\n<p>JAX \u8fd8\u6709\u4e00\u4e2a\u91cd\u8981\u7684\u7ec4\u5408\u51fd\u6570\u2014\u2014jit\uff0c\u5f00\u53d1\u8005\u53ef\u4ee5\u4f7f\u7528 jit transform \u5b9e\u73b0\u5373\u65f6\u7f16\u8bd1\u3002<\/p>\n<p><img decoding=\"async\" referrerpolicy=\"no-referrer\" src=\"https:\/\/segmentfault.com\/img\/bVc8F0Z\" alt=\"\" title=\"\" loading=\"lazy\"><\/p>\n<p>jit \u7ed3\u5408\u540e\u53f0\u53ef\u4ee5\u4f7f\u7528 XLA \u540e\u7aef\u7f16\u8bd1\u5668\u5c06\u64cd\u4f5c\u878d\u5408\u5728\u4e00\u8d77\uff0c\u6765\u81ea\u52a8\u5b9a\u4f4d CPU\u3001GPU \u548c TPU \u6216\u8005 ASIC\uff0c\u52a0\u901f\u8ba1\u7b97\u6a21\u578b\u548c\u7b97\u6cd5\u3002<\/p>\n<h3><strong>JAX \u52a8\u673a\u5256\u6790\uff1apmap<\/strong><\/h3>\n<p>\u6700\u540e\uff0c\u5982\u679c\u60f3\u5e76\u884c\u5316\u4f60\u7684\u4ee3\u7801\uff0c\u6709\u4e00\u4e2a\u548c vmap \u975e\u5e38\u76f8\u4f3c\u5f97\u8f6c\u6362\u53eb pmap\u3002<\/p>\n<p>\u901a\u8fc7\u4ee3\u7801\u8fd0\u884c pmap\uff0c\u5f00\u53d1\u8005\u80fd\u591f\u672c\u5730\u5b9a\u4f4d\u7cfb\u7edf\u4e2d\u7684\u591a\u4e2a\u5185\u6838\u6216\u4f60\u6709\u6743\u8bbf\u95ee\u7684 GPU\u3001TPU \u6216 ASIC \u96c6\u7fa4\u3002<\/p>\n<p><img decoding=\"async\" referrerpolicy=\"no-referrer\" src=\"https:\/\/segmentfault.com\/img\/bVc8F00\" alt=\"\" title=\"\" loading=\"lazy\"><\/p>\n<p>\u8fd9\u6700\u7ec8\u6210\u4e3a\u4e00\u4e2a\u975e\u5e38\u5f3a\u5927\u7684\u7cfb\u7edf\uff0c<strong>\u53ef\u4ee5\u5728\u6ca1\u6709\u592a\u591a\u989d\u5916\u4ee3\u7801\u7684\u60c5\u51b5\u4e0b\u6784\u5efa\u6211\u4eec\u7528\u7c7b\u4f3c\u4e8e NumPy \u7684\u719f\u6089 API\uff0c\u505a\u6df1\u5ea6\u5b66\u4e60\u7684\u5feb\u901f\u8ba1\u7b97\u7b49\u5de5\u4f5c\u8d1f\u8f7d<\/strong>\u3002<\/p>\n<p>JAX \u7684\u5173\u952e\u8bbe\u8ba1\u601d\u60f3<\/p>\n<p>\u901a\u8fc7\u4e0a\u8ff0\u5bf9\u6bd4\u53ef\u4ee5\u770b\u5230\uff0c JAX \u4e0d\u4ec5\u4e3a\u5f00\u53d1\u8005\u63d0\u4f9b\u4e86\u548c NumPy \u76f8\u4f3c\u7684 API\uff0c\u4e0a\u8ff0\u7684\u4e94\u5927\u51fd\u6570\u8f6c\u6362\u7ec4\u5408\u4e5f\u8ba9 JAX \u53ef\u4ee5\u5728\u4e0d\u9700\u8981\u989d\u5916\u4ee3\u7801\u7684\u60c5\u51b5\u4e0b\uff0c\u5e2e\u52a9\u5f00\u53d1\u8005\u6784\u5efa\u6df1\u5ea6\u5b66\u4e60\u5e94\u7528\u8fdb\u884c\u5feb\u901f\u8ba1\u7b97\u3002<\/p>\n<p>\u8fd9\u91cc\u7684\u5173\u952e\u601d\u60f3\u662f\uff1a<\/p>\n<p>1\uff09\u9996\u5148\uff0c\u5728 JAX \u4e2d\uff0cPython \u4ee3\u7801\u88ab\u8ffd\u6eaf\u5230\u4e2d\u95f4\u8868\u793a\uff0cJAX \u77e5\u9053\u5982\u4f55\u8f6c\u6362\u8fd9\u4e2a\u4e2d\u95f4\u8868\u793a\u3002<\/p>\n<p>2\uff09\u5728\u4e0b\u7bc7\u6587\u7ae0\u4e2d\u6211\u4eec\u4e5f\u5c06\u8be6\u7ec6\u5206\u6790 JAX \u7684\u5de5\u4f5c\u673a\u5236\uff1a\u540c\u6837\u7684\u4e2d\u95f4\u8868\u793a\uff0c\u901a\u8fc7\u5141\u8bb8 XLA \u8fdb\u884c\u7279\u5b9a\u9886\u57df (CPU\/GPU \u7b49) \u7684\u7f16\u8bd1\uff0c\u5982\u4f55\u6765\u7784\u51c6\u4e0d\u540c\u7684\u540e\u7aef\uff1b<\/p>\n<p>3\uff09\u53e6\u5916\uff0cJAX \u8fd8\u6709\u57fa\u4e8e NumPy \u548c SciPy \u7684\u9762\u5411\u7528\u6237\u7684 API\uff0c\u5982\u679c\u5f00\u53d1\u8005\u4e00\u76f4\u4f7f\u7528 Python \u7684\u6280\u672f\u6808\uff0c\u5e94\u8be5\u4f1a\u5bf9 JAX \u611f\u89c9\u76f8\u5f53\u719f\u6089\uff1b<\/p>\n<p>4\uff09\u6700\u540e\uff0cJAX \u63d0\u4f9b\u4e86<strong>\u529f\u80fd\u5f3a\u5927\u7684\u53d8\u6362\uff1agrad, git, vmap, pmap \u7b49\uff0c\u6765\u652f\u6301\u6df1\u5ea6\u5b66\u4e60\u7b49\u8ba1\u7b97<\/strong>\uff0c\u56e0\u6b64 JAX \u53ef\u4ee5\u505a\u5230\u4e4b\u524d NumPy \u4ee3\u7801\u65e0\u6cd5\u505a\u5230\u7684\u4e8b\u60c5\u3002<\/p>\n<p>\u901a\u8fc7\u524d\u9762\u7684\u4ecb\u7ecd\uff0c\u6211\u4eec\u53ef\u4ee5\u770b\u5230\uff0c\u5f00\u53d1\u8005\u719f\u6089\u7684 API \u548c\u8bed\u6cd5\u4ee5\u53ca\u56db\u79cd\u5f3a\u5927\u7684\u8f6c\u6362\u7ec4\u5408\u8ba9\u5f00\u53d1\u8005\u66f4\u52a0\u559c\u6b22 JAX\uff0c\u5e76\u8ba9\u6df1\u5ea6\u5b66\u4e60\u573a\u666f\u6216\u8005\u79d1\u5b66\u8ba1\u7b97\u53d8\u5f97\u975e\u5e38\u7b80\u4fbf\u3002<br \/>\u6b22\u8fce\u56de\u987e\u5173\u4e8e<a href=\"https:\/\/link.segmentfault.com\/?enc=81FsafVZMpwFBYZJA3GkNA%3D%3D.AcHWJC7m6y%2FltHf5cmodTw8MOUtLkhYbuhT%2BlmyKAYWqOgXHNe1FkBsg%2BC9USiV09EL32HrAtl5XYT6VILy82DuaImD%2FuR6DTaNvrn8ZHYxko6%2Fa4o%2FOKkZjcxoKE1YzSh8aFcbCo%2BrxpCQRRlapgoyTOvl57UObdpAtIvtk1irAfjXXRFbOitk57%2BHv04Dj\" rel=\"nofollow\">\u673a\u5668\u5b66\u4e60<\/a>\u7684\u5f80\u671f\u6587\u7ae0\uff0c\u4ee5\u53ca\u66f4\u591a\u9762\u5411\u5f00\u53d1\u8005\u7684\u6280\u672f\u5206\u4eab\u3002\u8bf7\u6301\u7eed\u5173\u6ce8 Build On Cloud \u5fae\u4fe1\u516c\u4f17\u53f7\uff01<\/p>\n<h3><strong>\u5f80\u671f\u63a8\u8350<\/strong><\/h3>\n<ul>\n<li><a href=\"https:\/\/link.segmentfault.com\/?enc=Gtt7Em4CeX%2FL9jGKKQaRLA%3D%3D.TNybEX5ihM%2FZTxeIEzNs%2F98XsPDplOPcLeUaddQ7gcq1qtajVotNi7%2BWNA0NNE2pT8rnw1bl5LxRDmMOExpika24Arn5PAqrt2AYRK4YLeEDMKwwd%2BMWJF6yqIt3idtd\" rel=\"nofollow\">\u673a\u5668\u5b66\u4e60\u6d1e\u5bdf | \u6316\u6398\u591a\u6a21\u6001\u6570\u636e\u673a\u5668\u5b66\u4e60\u7684\u4ef7\u503c<\/a><\/li>\n<li><a href=\"https:\/\/link.segmentfault.com\/?enc=rHcpFzl3Y%2Bodk%2F3BfQCX2w%3D%3D.YNvSUrdoCBh2rwVwnhb3BvUVKZnIlPk7BkpwdawC2GFS95mYkThwrNcUOZbJDqTxlEOlhxJQ%2F7xbRNQBYI6jko6cJXwnIY82JMc4xVZufO3qLUGMokC7Odidw6RmFwD%2B\" rel=\"nofollow\">\u673a\u5668\u5b66\u4e60\u6d1e\u5bdf | \u5206\u5e03\u5f0f\u8bad\u7ec3\u8ba9\u673a\u5668\u5b66\u4e60\u66f4\u52a0\u5feb\u901f\u51c6\u786e<\/a><\/li>\n<li><a href=\"https:\/\/link.segmentfault.com\/?enc=3O0yT2Axjc1k4QhJsNaT2A%3D%3D.HnmhUIjlZgg42eczQh8GrMG01NgsFjJf5SD76OrFYwomPdAwx6L01JOo%2Bbxgy0BMCjIO26oUhNLrCtFFCUL2lEEuTX7IrMmJnlJFPs%2FlfccrHj43Fu5XsiL9jgdSgwAn\" rel=\"nofollow\">\u673a\u5668\u5b66\u4e60\u6d1e\u5bdf | \u964d\u672c\u589e\u6548\uff0c\u65e0\u670d\u52a1\u5668\u63a8\u7406\u662f\u600e\u4e48\u505a\u5230\u7684\uff1f<\/a><\/li>\n<\/ul>\n<p><img decoding=\"async\" referrerpolicy=\"no-referrer\" src=\"https:\/\/segmentfault.com\/img\/bVc8F1w\" alt=\"\" title=\"\" loading=\"lazy\"><\/p>\n<p>\u4f5c\u8005<br \/>\u9ec4\u6d69\u6587<br \/>\u4e9a\u9a6c\u900a\u4e91\u79d1\u6280\u8d44\u6df1\u5f00\u53d1\u8005\u5e03\u9053\u5e08\uff0c\u4e13\u6ce8\u4e8e AI\/ML\u3001Data Science \u7b49\u3002\u62e5\u6709 20 \u591a\u5e74\u7535\u4fe1\u3001\u79fb\u52a8\u4e92\u8054\u7f51\u4ee5\u53ca\u4e91\u8ba1\u7b97\u7b49\u884c\u4e1a\u67b6\u6784\u8bbe\u8ba1\u3001\u6280\u672f\u53ca\u521b\u4e1a\u7ba1\u7406\u7b49\u4e30\u5bcc\u7ecf\u9a8c\uff0c\u66fe\u5c31\u804c\u4e8e Microsoft\u3001Sun Microsystems\u3001\u4e2d\u56fd\u7535\u4fe1\u7b49\u4f01\u4e1a\uff0c\u4e13\u6ce8\u4e3a\u6e38\u620f\u3001\u7535\u5546\u3001\u5a92\u4f53\u548c\u5e7f\u544a\u7b49\u4f01\u4e1a\u5ba2\u6237\u63d0\u4f9b AI\/ML\u3001\u6570\u636e\u5206\u6790\u548c\u4f01\u4e1a\u6570\u5b57\u5316\u8f6c\u578b\u7b49\u89e3\u51b3\u65b9\u6848\u54a8\u8be2\u670d\u52a1\u3002<\/p>\n<p>\u6587\u7ae0\u6765\u6e90\uff1a<a href=\"https:\/\/link.segmentfault.com\/?enc=fz5PuXst6DLmnIUb6D7AgQ%3D%3D.AMeFcZllQMlOGBUqAlcD5Xw0%2F2PvQyydjci1pw7IyI7ptuF50UT0Ze9BDEtB4TxMKXJqMj4m%2Fs%2FP%2B1W4M1JQXo9LLqxX7mxNsht4lB2tdazWAP02OCikc1Rz8mE39s0tlNcigw9wHbJEt1Lv%2B9wmYqd46CU63%2F11yJjcCSF37GiJ984WylOdXgdlc6pS9jZ0\" rel=\"nofollow\">https:\/\/dev.amazoncloud.cn\/column\/article\/63e33239e5e05b6ff89&#8230;<\/a><\/p>\n<p><a href=\"https:\/\/segmentfault.com\/a\/1190000043974830\">https:\/\/segmentfault.com\/a\/1190000043974830<\/a><\/p>\n","protected":false},"excerpt":{"rendered":"<p>\u5728\u4e4b\u524d\u7684\u300a\u673a\u5668\u5b66\u4e60\u6d1e\u5bdf\u300b\u7cfb\u5217\u6587\u7ae0\u4e2d\uff0c\u6211\u4eec\u5206\u522b\u9488\u5bf9\u4e8e\u591a\u6a21\u6001\u673a\u5668\u5b66\u4e60\u548c\u5206\u5e03\u5f0f\u8bad\u7ec3\u3001\u65e0\u670d\u52a1\u5668\u63a8\u7406\u8fdb\u884c\u4e86\u89e3\u8bfb\uff0c\u672c\u6587\u5c06\u4e3a [&hellip;]<\/p>\n","protected":false},"author":1,"featured_media":0,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"_jetpack_memberships_contains_paid_content":false,"footnotes":""},"categories":[4],"tags":[],"class_list":["post-10008044","post","type-post","status-publish","format-standard","hentry","category-share"],"jetpack_sharing_enabled":true,"jetpack_featured_media_url":"","_links":{"self":[{"href":"https:\/\/itmax.cn\/index.php?rest_route=\/wp\/v2\/posts\/10008044","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/itmax.cn\/index.php?rest_route=\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/itmax.cn\/index.php?rest_route=\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/itmax.cn\/index.php?rest_route=\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"https:\/\/itmax.cn\/index.php?rest_route=%2Fwp%2Fv2%2Fcomments&post=10008044"}],"version-history":[{"count":0,"href":"https:\/\/itmax.cn\/index.php?rest_route=\/wp\/v2\/posts\/10008044\/revisions"}],"wp:attachment":[{"href":"https:\/\/itmax.cn\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=10008044"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/itmax.cn\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=10008044"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/itmax.cn\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=10008044"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}