CPU device for tutorials

Reviewed By: nikhilaravi

Differential Revision: D22357376

fbshipit-source-id: c103f9b0c798d4425d642781b5bfbb1a27310270
This commit is contained in:
Jeremy Reizenstein
2020-07-03 08:51:19 -07:00
committed by Facebook GitHub Bot
parent 52979226bc
commit 275ddade66
5 changed files with 46 additions and 29 deletions

View File

@@ -114,7 +114,12 @@
"sys.path.append(os.path.abspath(''))\n",
"\n",
"# set for reproducibility\n",
"torch.manual_seed(42)"
"torch.manual_seed(42)\n",
"if torch.cuda.is_available():\n",
" device = torch.device(\"cuda:0\")\n",
"else:\n",
" device = torch.device(\"cpu\")\n",
" print(\"WARNING: CPU only, this will be slow!\")"
]
},
{
@@ -200,16 +205,16 @@
"\n",
"# create the relative cameras\n",
"cameras_relative = SfMPerspectiveCameras(\n",
" R = R_relative.cuda(),\n",
" T = T_relative.cuda(),\n",
" device = \"cuda\",\n",
" R = R_relative.to(device),\n",
" T = T_relative.to(device),\n",
" device = device,\n",
")\n",
"\n",
"# create the absolute ground truth cameras\n",
"cameras_absolute_gt = SfMPerspectiveCameras(\n",
" R = R_absolute_gt.cuda(),\n",
" T = T_absolute_gt.cuda(),\n",
" device = \"cuda\",\n",
" R = R_absolute_gt.to(device),\n",
" T = T_absolute_gt.to(device),\n",
" device = device,\n",
")\n",
"\n",
"# the number of absolute camera positions\n",
@@ -270,7 +275,7 @@
" SfMPerspectiveCameras(\n",
" R = cams.R[edges[:, i]],\n",
" T = cams.T[edges[:, i]],\n",
" device = \"cuda\",\n",
" device = device,\n",
" ).get_world_to_view_transform()\n",
" for i in (0, 1)\n",
" ]\n",
@@ -283,7 +288,7 @@
" cams_relative = SfMPerspectiveCameras(\n",
" R = matrix_rel[:, :3, :3],\n",
" T = matrix_rel[:, 3, :3],\n",
" device = \"cuda\",\n",
" device = device,\n",
" )\n",
" return cams_relative"
]
@@ -320,8 +325,8 @@
"outputs": [],
"source": [
"# initialize the absolute log-rotations/translations with random entries\n",
"log_R_absolute_init = torch.randn(N, 3).float().cuda()\n",
"T_absolute_init = torch.randn(N, 3).float().cuda()\n",
"log_R_absolute_init = torch.randn(N, 3, dtype=torch.float32, device=device)\n",
"T_absolute_init = torch.randn(N, 3, dtype=torch.float32, device=device)\n",
"\n",
"# furthermore, we know that the first camera is a trivial one \n",
"# (see the description above)\n",
@@ -337,7 +342,7 @@
"# the mask the specifies which cameras are going to be optimized\n",
"# (since we know the first camera is already correct, \n",
"# we only optimize over the 2nd-to-last cameras)\n",
"camera_mask = torch.ones(N, 1).float().cuda()\n",
"camera_mask = torch.ones(N, 1, dtype=torch.float32, device=device)\n",
"camera_mask[0] = 0.\n",
"\n",
"# init the optimizer\n",
@@ -358,7 +363,7 @@
" cameras_absolute = SfMPerspectiveCameras(\n",
" R = R_absolute,\n",
" T = T_absolute * camera_mask,\n",
" device = \"cuda\",\n",
" device = device,\n",
" )\n",
"\n",
" # compute the relative cameras as a compositon of the absolute cameras\n",