使用 Colaboratory 创建

Author: williamyang1991

Diff for: notebooks/inference_playground.ipynb

@@ -327,7 +327,7 @@
         "opts.device = device\n",
         "encoder = pSp(opts)\n",
         "encoder.eval()\n",
-        "encoder.cuda()\n",
+        "encoder = encoder.to(device)\n",
         "\n",
         "# load extrinsic style code\n",
         "exstyles = np.load(os.path.join(MODEL_DIR, style_type, MODEL_PATHS[style_type+'-S'][\"name\"]), allow_pickle='TRUE').item()\n",
@@ -462,9 +462,9 @@
       "outputs": [],
       "source": [
         "if if_align_face:\n",
-        "    I = transform(run_alignment(image_path)).unsqueeze(dim=0).cuda()\n",
+        "    I = transform(run_alignment(image_path)).unsqueeze(dim=0).to(device)\n",
         "else:\n",
-        "    I = F.adaptive_avg_pool2d(load_image(image_path).cuda(), 256)"
+        "    I = F.adaptive_avg_pool2d(load_image(image_path).to(device), 256)"
       ]
     },
     {
@@ -612,7 +612,7 @@
         "                            z_plus_latent=True, return_z_plus_latent=True, resize=False)    \n",
         "    img_rec = torch.clamp(img_rec.detach(), -1, 1)\n",
         "    \n",
-        "    latent = torch.tensor(exstyles[stylename]).repeat(2,1,1).cuda()\n",
+        "    latent = torch.tensor(exstyles[stylename]).repeat(2,1,1).to(device)\n",
         "    # latent[0] for both color and structrue transfer and latent[1] for only structrue transfer\n",
         "    latent[1,7:18] = instyle[0,7:18]\n",
         "    exstyle = generator.generator.style(latent.reshape(latent.shape[0]*latent.shape[1], latent.shape[2])).reshape(latent.shape)\n",
@@ -824,9 +824,9 @@
       ],
       "source": [
         "with torch.no_grad():\n",
-        "    latent = torch.tensor(exstyles[stylename]).repeat(6,1,1).cuda()\n",
-        "    latent2 = torch.tensor(exstyles[stylename2]).repeat(6,1,1).cuda()\n",
-        "    fuse_weight = torch.arange(6).reshape(6,1,1).cuda() / 5.0\n",
+        "    latent = torch.tensor(exstyles[stylename]).repeat(6,1,1).to(device)\n",
+        "    latent2 = torch.tensor(exstyles[stylename2]).repeat(6,1,1).to(device)\n",
+        "    fuse_weight = torch.arange(6).reshape(6,1,1).to(device) / 5.0\n",
         "    fuse_latent = latent * fuse_weight + latent2 * (1-fuse_weight)\n",
         "    exstyle = generator.generator.style(fuse_latent.reshape(fuse_latent.shape[0]*fuse_latent.shape[1], fuse_latent.shape[2])).reshape(fuse_latent.shape)\n",
         "    \n",
@@ -871,11 +871,11 @@
         "batch = 6 # sample 6 style codes\n",
         "\n",
         "with torch.no_grad():\n",
-        "    instyle = torch.randn(6, 512).cuda()\n",
+        "    instyle = torch.randn(6, 512).to(device)\n",
         "    # sample structure codes\n",
-        "    res_in = icpts.icp.netT(torch.randn(batch, 128).cuda()).reshape(-1,7,512)\n",
+        "    res_in = icpts.icp.netT(torch.randn(batch, 128).to(device)).reshape(-1,7,512)\n",
         "    # sample color codes\n",
-        "    ada_in = icptc.icp.netT(torch.randn(batch, 128).cuda()).reshape(-1,11,512)\n",
+        "    ada_in = icptc.icp.netT(torch.randn(batch, 128).to(device)).reshape(-1,11,512)\n",
         "\n",
         "    # concatenate two codes to form the complete extrinsic style code\n",
         "    latent = torch.cat((res_in, ada_in), dim=1)\n",