Skip to content

Instantly share code, notes, and snippets.

@fehiepsi fehiepsi/cpu_vs_gpu.ipynb
Last active May 30, 2019

Embed
What would you like to do?
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import jax.numpy as np\n",
"from jax import jit, grad, lax\n",
"from jax.config import config; config.update(\"jax_platform_name\", \"gpu\")\n",
"\n",
"import numpyro.distributions as dist\n",
"from numpyro.examples.datasets import SP500, load_dataset\n",
"from numpyro.handlers import sample\n",
"from numpyro.hmc_util import initialize_model\n",
"from numpyro.mcmc import hmc"
]
},
{
"cell_type": "code",
"execution_count": 23,
"metadata": {},
"outputs": [],
"source": [
"_, fetch = load_dataset(SP500, shuffle=False)\n",
"dates, returns = fetch()"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [],
"source": [
"def potential_fn(x):\n",
" return ((x - returns) ** 2).mean()\n",
"\n",
"@jit\n",
"def loop(x):\n",
" def body_fn(i, x):\n",
" return x + 0.01 * grad(potential_fn)(x)\n",
"\n",
" return lax.fori_loop(0, 1000, body_fn, x)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"36.1 ms ± 2.84 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
]
}
],
"source": [
"%%timeit\n",
"loop(np.array(0.))"
]
},
{
"cell_type": "code",
"execution_count": 24,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"DeviceArray(-9989383., dtype=float32)"
]
},
"execution_count": 24,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"loop(np.array(0.))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here is the result in CPU."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"161 µs ± 1.81 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
]
}
],
"source": [
"%%timeit\n",
"loop(np.array(0.))"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"def loop(x):\n",
" def body_fn(i, x):\n",
" y = potential_fn(x)\n",
" return x + 0.01 * torch.autograd.grad(y, x)[0]\n",
" \n",
" for i in range(1000):\n",
" x = body_fn(i, x)\n",
" return x"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [],
"source": [
"returns = torch.from_numpy(returns).float()\n",
"x = torch.tensor(0., requires_grad=True)"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"344 ms ± 9.17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
]
}
],
"source": [
"%%timeit\n",
"loop(x)"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"tensor(-9989380., grad_fn=<AddBackward0>)"
]
},
"execution_count": 20,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"loop(x)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python (pydata)",
"language": "python",
"name": "pydata"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.6.8"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.