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",

View File

@ -84,7 +84,7 @@
"import os\n",
"import torch\n",
"import numpy as np\n",
"from tqdm import tqdm_notebook\n",
"from tqdm.notebook import tqdm\n",
"import imageio\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
@ -159,8 +159,11 @@
"outputs": [],
"source": [
"# Set the cuda device \n",
"device = torch.device(\"cuda:0\")\n",
"torch.cuda.set_device(device)\n",
"if torch.cuda.is_available():\n",
" device = torch.device(\"cuda:0\")\n",
" torch.cuda.set_device(device)\n",
"else:\n",
" device = torch.device(\"cpu\")\n",
"\n",
"# Load the obj and ignore the textures and materials.\n",
"verts, faces_idx, _ = load_obj(\"./data/teapot.obj\")\n",
@ -433,7 +436,7 @@
"plt.subplot(1, 2, 2)\n",
"plt.imshow(model.image_ref.cpu().numpy().squeeze())\n",
"plt.grid(False)\n",
"plt.title(\"Reference silhouette\")\n"
"plt.title(\"Reference silhouette\");\n"
]
},
{
@ -472,7 +475,7 @@
},
"outputs": [],
"source": [
"loop = tqdm_notebook(range(200))\n",
"loop = tqdm(range(200))\n",
"for i in loop:\n",
" optimizer.zero_grad()\n",
" loss, _ = model()\n",

View File

@ -108,7 +108,7 @@
" mesh_normal_consistency,\n",
")\n",
"import numpy as np\n",
"from tqdm import tqdm_notebook\n",
"from tqdm.notebook import tqdm\n",
"%matplotlib notebook \n",
"from mpl_toolkits.mplot3d import Axes3D\n",
"import matplotlib.pyplot as plt\n",
@ -117,7 +117,11 @@
"mpl.rcParams['figure.dpi'] = 80\n",
"\n",
"# Set the device\n",
"device = torch.device(\"cuda:0\")"
"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!\")"
]
},
{
@ -342,7 +346,7 @@
"w_laplacian = 0.1 \n",
"# Plot period for the losses\n",
"plot_period = 250\n",
"loop = tqdm_notebook(range(Niter))\n",
"loop = tqdm(range(Niter))\n",
"\n",
"chamfer_losses = []\n",
"laplacian_losses = []\n",
@ -428,7 +432,7 @@
"ax.legend(fontsize=\"16\")\n",
"ax.set_xlabel(\"Iteration\", fontsize=\"16\")\n",
"ax.set_ylabel(\"Loss\", fontsize=\"16\")\n",
"ax.set_title(\"Loss vs iterations\", fontsize=\"16\")"
"ax.set_title(\"Loss vs iterations\", fontsize=\"16\");"
]
},
{

View File

@ -52,7 +52,6 @@
"outputs": [],
"source": [
"import os\n",
"os.chdir('../..')\n",
"import torch\n",
"import torch.nn.functional as F\n",
"import matplotlib.pyplot as plt\n",
@ -110,8 +109,11 @@
"outputs": [],
"source": [
"# Setup\n",
"device = torch.device(\"cuda:0\")\n",
"torch.cuda.set_device(device)\n",
"if torch.cuda.is_available():\n",
" device = torch.device(\"cuda:0\")\n",
" torch.cuda.set_device(device)\n",
"else:\n",
" device = torch.device(\"cpu\")\n",
"\n",
"# Set paths\n",
"DATA_DIR = \"./data\"\n",
@ -177,7 +179,7 @@
"plt.figure(figsize=(10, 10))\n",
"plt.imshow(images[0, ..., :3].cpu().numpy())\n",
"plt.grid(\"off\")\n",
"plt.axis(\"off\")"
"plt.axis(\"off\");"
]
},
{
@ -225,7 +227,7 @@
"plt.figure(figsize=(10, 10))\n",
"plt.imshow(images[0, ..., :3].cpu().numpy())\n",
"plt.grid(\"off\")\n",
"plt.axis(\"off\")"
"plt.axis(\"off\");"
]
}
],

View File

@ -217,8 +217,11 @@
"outputs": [],
"source": [
"# Setup\n",
"device = torch.device(\"cuda:0\")\n",
"torch.cuda.set_device(device)\n",
"if torch.cuda.is_available():\n",
" device = torch.device(\"cuda:0\")\n",
" torch.cuda.set_device(device)\n",
"else:\n",
" device = torch.device(\"cpu\")\n",
"\n",
"# Set paths\n",
"DATA_DIR = \"./data\"\n",