使用 Colaboratory 创建

Author: williamyang1991

Diff for: notebooks/inference_playground.ipynb

@@ -312,17 +312,19 @@
         ")\n",
         "\n",
         "# load DualStyleGAN\n",
-        "generator = DualStyleGAN(1024, 512, 8, 2, res_index=6).cuda()\n",
+        "generator = DualStyleGAN(1024, 512, 8, 2, res_index=6)\n",
         "generator.eval()\n",
-        "ckpt = torch.load(os.path.join(MODEL_DIR, style_type, 'generator.pt'))\n",
+        "ckpt = torch.load(os.path.join(MODEL_DIR, style_type, 'generator.pt'), map_location=lambda storage, loc: storage)\n",
         "generator.load_state_dict(ckpt[\"g_ema\"])\n",
+        "generator = generator.to(device)\n",
         "\n",
         "# load encoder\n",
         "model_path = os.path.join(MODEL_DIR, 'encoder.pt')\n",
         "ckpt = torch.load(model_path, map_location='cpu')\n",
         "opts = ckpt['opts']\n",
         "opts['checkpoint_path'] = model_path\n",
         "opts = Namespace(**opts)\n",
+        "opts.device = device\n",
         "encoder = pSp(opts)\n",
         "encoder.eval()\n",
         "encoder.cuda()\n",
@@ -333,9 +335,11 @@
         "# load sampler network\n",
         "icptc = ICPTrainer(np.empty([0,512*11]), 128)\n",
         "icpts = ICPTrainer(np.empty([0,512*7]), 128)\n",
-        "ckpt = torch.load(os.path.join(MODEL_DIR, style_type, 'sampler.pt'))\n",
+        "ckpt = torch.load(os.path.join(MODEL_DIR, style_type, 'sampler.pt'), map_location=lambda storage, loc: storage)\n",
         "icptc.icp.netT.load_state_dict(ckpt['color'])\n",
         "icpts.icp.netT.load_state_dict(ckpt['structure'])\n",
+        "icptc.icp.netT = icptc.icp.netT.to(device)\n",
+        "icpts.icp.netT = icpts.icp.netT.to(device)\n",
         "\n",
         "print('Model successfully loaded!')"
       ]