Skip to content

Instantly share code, notes, and snippets.

@ferrine
Created January 27, 2019 20:46
Show Gist options
  • Save ferrine/ac287080f8fc28d40b7823589fbd06cc to your computer and use it in GitHub Desktop.
Save ferrine/ac287080f8fc28d40b7823589fbd06cc to your computer and use it in GitHub Desktop.
geoopt-stiefel-benchmark
Display the source blob
Display the rendered blob
Raw
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"0.4.1.post2\n"
]
}
],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"print(torch.__version__)"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Sun Jan 27 23:36:09 2019 \n",
"+-----------------------------------------------------------------------------+\n",
"| NVIDIA-SMI 410.48 Driver Version: 410.48 |\n",
"|-------------------------------+----------------------+----------------------+\n",
"| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |\n",
"| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |\n",
"|===============================+======================+======================|\n",
"| 0 GeForce GTX 108... Off | 00000000:0A:00.0 Off | N/A |\n",
"| 28% 46C P0 54W / 250W | 0MiB / 11178MiB | 0% Default |\n",
"+-------------------------------+----------------------+----------------------+\n",
"| 1 GeForce GTX 108... Off | 00000000:41:00.0 Off | N/A |\n",
"| 23% 42C P0 53W / 250W | 0MiB / 11177MiB | 0% Default |\n",
"+-------------------------------+----------------------+----------------------+\n",
" \n",
"+-----------------------------------------------------------------------------+\n",
"| Processes: GPU Memory |\n",
"| GPU PID Type Process name Usage |\n",
"|=============================================================================|\n",
"| No running processes found |\n",
"+-----------------------------------------------------------------------------+\n"
]
}
],
"source": [
"!nvidia-smi"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def retr(x, u, t):\n",
" a = u @ x.transpose(-1, -2) - x @ u.transpose(-1, -2)\n",
" rhs = u + t / 2 * a @ u\n",
" lhs = -t / 2 * a\n",
" lhs[..., torch.arange(a.shape[-2]), torch.arange(x.shape[-2])] += 1\n",
" qv, _ = torch.gesv(rhs, lhs)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"def proju(x, u):\n",
" return u - x @ u.transpose(-1, -2) @ x"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"def timeit(cuda, n):\n",
" Q = torch.qr(torch.randn(n, n))[0]\n",
" U = proju(Q, torch.randn(n, n))\n",
" if cuda:\n",
" Q = Q.cuda()\n",
" U = U.cuda()\n",
" %timeit retr(Q, U, 1.)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2.26 ms ± 21.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
]
}
],
"source": [
"timeit(True, 10)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2.48 ms ± 54.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
]
}
],
"source": [
"timeit(True, 100)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"12.1 ms ± 122 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
]
}
],
"source": [
"timeit(True, 1000)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2.39 s ± 19.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
]
}
],
"source": [
"timeit(True, 10000)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"69.5 µs ± 121 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
]
}
],
"source": [
"timeit(False, 10)"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"358 µs ± 3.96 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
]
}
],
"source": [
"timeit(False, 100)"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"48.2 ms ± 1.94 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
]
}
],
"source": [
"timeit(False, 1000)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#timeit(False, 10000)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "pytorch-0.4.1",
"language": "python",
"name": "pytorch-0.4.1"
},
"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.7.2"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Display the source blob
Display the rendered blob
Raw
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment