0728-1750

diffusion.py

@@ -560,7 +560,7 @@ class DDPM21CM:
     def save(self, ep):
         # save model
         # if self.accelerator.is_main_process:
-        if torch.cuda.current_device() == 0:
+        if self.config.global_rank == 0:# and torch.cuda.current_device() == 0:
             if ep == self.config.n_epoch-1 or (ep+1) % self.config.save_period == 0:
                 self.nn_model.eval()
                 with torch.no_grad():
@@ -674,7 +674,7 @@ def train(rank, world_size, local_world_size, master_addr, master_port):
     config = TrainConfig()
     config.device = f"cuda:{rank}"
     config.world_size = local_world_size
-    
+    config.global_rank = global_rank 
     #[3200]#[200]#[1600,3200,6400,12800,25600]
     #for i, num_image in enumerate(num_train_image_list):
         #config.num_image = num_image

quantify_results.ipynb

@@ -76,121 +76,36 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "total 1187288\n",
-      "drwxr-xr-x 163 bxia34     12288 Jul 28 17:27 \u001b[0m\u001b[01;34mlogs\u001b[0m/\n",
-      "-rw-r--r--   1 bxia34 607850537 Jul 28 17:28 model_state-N30-device_count3-epoch4-172.27.145.66\n",
-      "-rw-r--r--   1 bxia34 607850537 Jul 28 17:28 model_state-N30-device_count3-epoch4-172.27.145.67\n"
+      "total 1187284\n",
+      "drwxr-xr-x 164 bxia34     12288 Jul 28 17:34 \u001b[0m\u001b[01;34mlogs\u001b[0m/\n",
+      "-rw-r--r--   1 bxia34 607850537 Jul 28 17:35 model_state-N30-device_count3-epoch4-172.27.145.66\n",
+      "-rw-r--r--   1 bxia34 607850537 Jul 28 17:35 model_state-N30-device_count3-epoch4-172.27.145.67\n"
      ]
     }
    ],
    "source": [
-    "ll outputs"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "ll outputs\n",
     "model0 = torch.load(\"outputs/model_state-N30-device_count3-epoch4-172.27.145.66\")\n",
-    "model1 = torch.load(\"outputs/model_state-N30-device_count3-epoch4-172.27.145.67\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "dict_keys(['epoch', 'unet_state_dict'])"
-      ]
-     },
-     "execution_count": 15,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model0.keys()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "4"
-      ]
-     },
-     "execution_count": 16,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model0['epoch']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 17,
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "model1 = torch.load(\"outputs/model_state-N30-device_count3-epoch4-172.27.145.67\")\n",
+    "model00 = torch.load(\"outputs/model_state-N30-device_count3-epoch4-172.27.145.66\")\n",
+    "model11 = torch.load(\"outputs/model_state-N30-device_count3-epoch4-172.27.145.67\")\n",
+    "model0.keys()\n",
+    "model0['epoch']\n",
     "state0 = model0['unet_state_dict']\n",
-    "state1 = model1['unet_state_dict']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "odict_keys(['token_embedding.weight', 'token_embedding.bias', 'time_embed.0.weight', 'time_embed.0.bias', 'time_embed.2.weight', 'time_embed.2.bias', 'input_blocks.0.0.weight', 'input_blocks.0.0.bias', 'input_blocks.1.0.in_layers.0.weight', 'input_blocks.1.0.in_layers.0.bias', 'input_blocks.1.0.in_layers.2.weight', 'input_blocks.1.0.in_layers.2.bias', 'input_blocks.1.0.emb_layers.1.weight', 'input_blocks.1.0.emb_layers.1.bias', 'input_blocks.1.0.out_layers.0.weight', 'input_blocks.1.0.out_layers.0.bias', 'input_blocks.1.0.out_layers.3.weight', 'input_blocks.1.0.out_layers.3.bias', 'input_blocks.2.0.in_layers.0.weight', 'input_blocks.2.0.in_layers.0.bias', 'input_blocks.2.0.in_layers.2.weight', 'input_blocks.2.0.in_layers.2.bias', 'input_blocks.2.0.emb_layers.1.weight', 'input_blocks.2.0.emb_layers.1.bias', 'input_blocks.2.0.out_layers.0.weight', 'input_blocks.2.0.out_layers.0.bias', 'input_blocks.2.0.out_layers.3.weight', 'input_blocks.2.0.out_layers.3.bias', 'input_blocks.3.0.op.weight', 'input_blocks.3.0.op.bias', 'input_blocks.4.0.in_layers.0.weight', 'input_blocks.4.0.in_layers.0.bias', 'input_blocks.4.0.in_layers.2.weight', 'input_blocks.4.0.in_layers.2.bias', 'input_blocks.4.0.emb_layers.1.weight', 'input_blocks.4.0.emb_layers.1.bias', 'input_blocks.4.0.out_layers.0.weight', 'input_blocks.4.0.out_layers.0.bias', 'input_blocks.4.0.out_layers.3.weight', 'input_blocks.4.0.out_layers.3.bias', 'input_blocks.5.0.in_layers.0.weight', 'input_blocks.5.0.in_layers.0.bias', 'input_blocks.5.0.in_layers.2.weight', 'input_blocks.5.0.in_layers.2.bias', 'input_blocks.5.0.emb_layers.1.weight', 'input_blocks.5.0.emb_layers.1.bias', 'input_blocks.5.0.out_layers.0.weight', 'input_blocks.5.0.out_layers.0.bias', 'input_blocks.5.0.out_layers.3.weight', 'input_blocks.5.0.out_layers.3.bias', 'input_blocks.6.0.op.weight', 'input_blocks.6.0.op.bias', 'input_blocks.7.0.in_layers.0.weight', 'input_blocks.7.0.in_layers.0.bias', 'input_blocks.7.0.in_layers.2.weight', 'input_blocks.7.0.in_layers.2.bias', 'input_blocks.7.0.emb_layers.1.weight', 'input_blocks.7.0.emb_layers.1.bias', 'input_blocks.7.0.out_layers.0.weight', 'input_blocks.7.0.out_layers.0.bias', 'input_blocks.7.0.out_layers.3.weight', 'input_blocks.7.0.out_layers.3.bias', 'input_blocks.7.0.skip_connection.weight', 'input_blocks.7.0.skip_connection.bias', 'input_blocks.7.1.norm.weight', 'input_blocks.7.1.norm.bias', 'input_blocks.7.1.qkv.weight', 'input_blocks.7.1.qkv.bias', 'input_blocks.7.1.proj_out.weight', 'input_blocks.7.1.proj_out.bias', 'input_blocks.8.0.in_layers.0.weight', 'input_blocks.8.0.in_layers.0.bias', 'input_blocks.8.0.in_layers.2.weight', 'input_blocks.8.0.in_layers.2.bias', 'input_blocks.8.0.emb_layers.1.weight', 'input_blocks.8.0.emb_layers.1.bias', 'input_blocks.8.0.out_layers.0.weight', 'input_blocks.8.0.out_layers.0.bias', 'input_blocks.8.0.out_layers.3.weight', 'input_blocks.8.0.out_layers.3.bias', 'input_blocks.8.1.norm.weight', 'input_blocks.8.1.norm.bias', 'input_blocks.8.1.qkv.weight', 'input_blocks.8.1.qkv.bias', 'input_blocks.8.1.proj_out.weight', 'input_blocks.8.1.proj_out.bias', 'input_blocks.9.0.op.weight', 'input_blocks.9.0.op.bias', 'input_blocks.10.0.in_layers.0.weight', 'input_blocks.10.0.in_layers.0.bias', 'input_blocks.10.0.in_layers.2.weight', 'input_blocks.10.0.in_layers.2.bias', 'input_blocks.10.0.emb_layers.1.weight', 'input_blocks.10.0.emb_layers.1.bias', 'input_blocks.10.0.out_layers.0.weight', 'input_blocks.10.0.out_layers.0.bias', 'input_blocks.10.0.out_layers.3.weight', 'input_blocks.10.0.out_layers.3.bias', 'input_blocks.10.1.norm.weight', 'input_blocks.10.1.norm.bias', 'input_blocks.10.1.qkv.weight', 'input_blocks.10.1.qkv.bias', 'input_blocks.10.1.proj_out.weight', 'input_blocks.10.1.proj_out.bias', 'input_blocks.11.0.in_layers.0.weight', 'input_blocks.11.0.in_layers.0.bias', 'input_blocks.11.0.in_layers.2.weight', 'input_blocks.11.0.in_layers.2.bias', 'input_blocks.11.0.emb_layers.1.weight', 'input_blocks.11.0.emb_layers.1.bias', 'input_blocks.11.0.out_layers.0.weight', 'input_blocks.11.0.out_layers.0.bias', 'input_blocks.11.0.out_layers.3.weight', 'input_blocks.11.0.out_layers.3.bias', 'input_blocks.11.1.norm.weight', 'input_blocks.11.1.norm.bias', 'input_blocks.11.1.qkv.weight', 'input_blocks.11.1.qkv.bias', 'input_blocks.11.1.proj_out.weight', 'input_blocks.11.1.proj_out.bias', 'input_blocks.12.0.op.weight', 'input_blocks.12.0.op.bias', 'input_blocks.13.0.in_layers.0.weight', 'input_blocks.13.0.in_layers.0.bias', 'input_blocks.13.0.in_layers.2.weight', 'input_blocks.13.0.in_layers.2.bias', 'input_blocks.13.0.emb_layers.1.weight', 'input_blocks.13.0.emb_layers.1.bias', 'input_blocks.13.0.out_layers.0.weight', 'input_blocks.13.0.out_layers.0.bias', 'input_blocks.13.0.out_layers.3.weight', 'input_blocks.13.0.out_layers.3.bias', 'input_blocks.14.0.in_layers.0.weight', 'input_blocks.14.0.in_layers.0.bias', 'input_blocks.14.0.in_layers.2.weight', 'input_blocks.14.0.in_layers.2.bias', 'input_blocks.14.0.emb_layers.1.weight', 'input_blocks.14.0.emb_layers.1.bias', 'input_blocks.14.0.out_layers.0.weight', 'input_blocks.14.0.out_layers.0.bias', 'input_blocks.14.0.out_layers.3.weight', 'input_blocks.14.0.out_layers.3.bias', 'middle_block.0.in_layers.0.weight', 'middle_block.0.in_layers.0.bias', 'middle_block.0.in_layers.2.weight', 'middle_block.0.in_layers.2.bias', 'middle_block.0.emb_layers.1.weight', 'middle_block.0.emb_layers.1.bias', 'middle_block.0.out_layers.0.weight', 'middle_block.0.out_layers.0.bias', 'middle_block.0.out_layers.3.weight', 'middle_block.0.out_layers.3.bias', 'middle_block.1.norm.weight', 'middle_block.1.norm.bias', 'middle_block.1.qkv.weight', 'middle_block.1.qkv.bias', 'middle_block.1.proj_out.weight', 'middle_block.1.proj_out.bias', 'middle_block.2.in_layers.0.weight', 'middle_block.2.in_layers.0.bias', 'middle_block.2.in_layers.2.weight', 'middle_block.2.in_layers.2.bias', 'middle_block.2.emb_layers.1.weight', 'middle_block.2.emb_layers.1.bias', 'middle_block.2.out_layers.0.weight', 'middle_block.2.out_layers.0.bias', 'middle_block.2.out_layers.3.weight', 'middle_block.2.out_layers.3.bias', 'output_blocks.0.0.in_layers.0.weight', 'output_blocks.0.0.in_layers.0.bias', 'output_blocks.0.0.in_layers.2.weight', 'output_blocks.0.0.in_layers.2.bias', 'output_blocks.0.0.emb_layers.1.weight', 'output_blocks.0.0.emb_layers.1.bias', 'output_blocks.0.0.out_layers.0.weight', 'output_blocks.0.0.out_layers.0.bias', 'output_blocks.0.0.out_layers.3.weight', 'output_blocks.0.0.out_layers.3.bias', 'output_blocks.0.0.skip_connection.weight', 'output_blocks.0.0.skip_connection.bias', 'output_blocks.1.0.in_layers.0.weight', 'output_blocks.1.0.in_layers.0.bias', 'output_blocks.1.0.in_layers.2.weight', 'output_blocks.1.0.in_layers.2.bias', 'output_blocks.1.0.emb_layers.1.weight', 'output_blocks.1.0.emb_layers.1.bias', 'output_blocks.1.0.out_layers.0.weight', 'output_blocks.1.0.out_layers.0.bias', 'output_blocks.1.0.out_layers.3.weight', 'output_blocks.1.0.out_layers.3.bias', 'output_blocks.1.0.skip_connection.weight', 'output_blocks.1.0.skip_connection.bias', 'output_blocks.2.0.in_layers.0.weight', 'output_blocks.2.0.in_layers.0.bias', 'output_blocks.2.0.in_layers.2.weight', 'output_blocks.2.0.in_layers.2.bias', 'output_blocks.2.0.emb_layers.1.weight', 'output_blocks.2.0.emb_layers.1.bias', 'output_blocks.2.0.out_layers.0.weight', 'output_blocks.2.0.out_layers.0.bias', 'output_blocks.2.0.out_layers.3.weight', 'output_blocks.2.0.out_layers.3.bias', 'output_blocks.2.0.skip_connection.weight', 'output_blocks.2.0.skip_connection.bias', 'output_blocks.2.1.conv.weight', 'output_blocks.2.1.conv.bias', 'output_blocks.3.0.in_layers.0.weight', 'output_blocks.3.0.in_layers.0.bias', 'output_blocks.3.0.in_layers.2.weight', 'output_blocks.3.0.in_layers.2.bias', 'output_blocks.3.0.emb_layers.1.weight', 'output_blocks.3.0.emb_layers.1.bias', 'output_blocks.3.0.out_layers.0.weight', 'output_blocks.3.0.out_layers.0.bias', 'output_blocks.3.0.out_layers.3.weight', 'output_blocks.3.0.out_layers.3.bias', 'output_blocks.3.0.skip_connection.weight', 'output_blocks.3.0.skip_connection.bias', 'output_blocks.3.1.norm.weight', 'output_blocks.3.1.norm.bias', 'output_blocks.3.1.qkv.weight', 'output_blocks.3.1.qkv.bias', 'output_blocks.3.1.proj_out.weight', 'output_blocks.3.1.proj_out.bias', 'output_blocks.4.0.in_layers.0.weight', 'output_blocks.4.0.in_layers.0.bias', 'output_blocks.4.0.in_layers.2.weight', 'output_blocks.4.0.in_layers.2.bias', 'output_blocks.4.0.emb_layers.1.weight', 'output_blocks.4.0.emb_layers.1.bias', 'output_blocks.4.0.out_layers.0.weight', 'output_blocks.4.0.out_layers.0.bias', 'output_blocks.4.0.out_layers.3.weight', 'output_blocks.4.0.out_layers.3.bias', 'output_blocks.4.0.skip_connection.weight', 'output_blocks.4.0.skip_connection.bias', 'output_blocks.4.1.norm.weight', 'output_blocks.4.1.norm.bias', 'output_blocks.4.1.qkv.weight', 'output_blocks.4.1.qkv.bias', 'output_blocks.4.1.proj_out.weight', 'output_blocks.4.1.proj_out.bias', 'output_blocks.5.0.in_layers.0.weight', 'output_blocks.5.0.in_layers.0.bias', 'output_blocks.5.0.in_layers.2.weight', 'output_blocks.5.0.in_layers.2.bias', 'output_blocks.5.0.emb_layers.1.weight', 'output_blocks.5.0.emb_layers.1.bias', 'output_blocks.5.0.out_layers.0.weight', 'output_blocks.5.0.out_layers.0.bias', 'output_blocks.5.0.out_layers.3.weight', 'output_blocks.5.0.out_layers.3.bias', 'output_blocks.5.0.skip_connection.weight', 'output_blocks.5.0.skip_connection.bias', 'output_blocks.5.1.norm.weight', 'output_blocks.5.1.norm.bias', 'output_blocks.5.1.qkv.weight', 'output_blocks.5.1.qkv.bias', 'output_blocks.5.1.proj_out.weight', 'output_blocks.5.1.proj_out.bias', 'output_blocks.5.2.conv.weight', 'output_blocks.5.2.conv.bias', 'output_blocks.6.0.in_layers.0.weight', 'output_blocks.6.0.in_layers.0.bias', 'output_blocks.6.0.in_layers.2.weight', 'output_blocks.6.0.in_layers.2.bias', 'output_blocks.6.0.emb_layers.1.weight', 'output_blocks.6.0.emb_layers.1.bias', 'output_blocks.6.0.out_layers.0.weight', 'output_blocks.6.0.out_layers.0.bias', 'output_blocks.6.0.out_layers.3.weight', 'output_blocks.6.0.out_layers.3.bias', 'output_blocks.6.0.skip_connection.weight', 'output_blocks.6.0.skip_connection.bias', 'output_blocks.6.1.norm.weight', 'output_blocks.6.1.norm.bias', 'output_blocks.6.1.qkv.weight', 'output_blocks.6.1.qkv.bias', 'output_blocks.6.1.proj_out.weight', 'output_blocks.6.1.proj_out.bias', 'output_blocks.7.0.in_layers.0.weight', 'output_blocks.7.0.in_layers.0.bias', 'output_blocks.7.0.in_layers.2.weight', 'output_blocks.7.0.in_layers.2.bias', 'output_blocks.7.0.emb_layers.1.weight', 'output_blocks.7.0.emb_layers.1.bias', 'output_blocks.7.0.out_layers.0.weight', 'output_blocks.7.0.out_layers.0.bias', 'output_blocks.7.0.out_layers.3.weight', 'output_blocks.7.0.out_layers.3.bias', 'output_blocks.7.0.skip_connection.weight', 'output_blocks.7.0.skip_connection.bias', 'output_blocks.7.1.norm.weight', 'output_blocks.7.1.norm.bias', 'output_blocks.7.1.qkv.weight', 'output_blocks.7.1.qkv.bias', 'output_blocks.7.1.proj_out.weight', 'output_blocks.7.1.proj_out.bias', 'output_blocks.8.0.in_layers.0.weight', 'output_blocks.8.0.in_layers.0.bias', 'output_blocks.8.0.in_layers.2.weight', 'output_blocks.8.0.in_layers.2.bias', 'output_blocks.8.0.emb_layers.1.weight', 'output_blocks.8.0.emb_layers.1.bias', 'output_blocks.8.0.out_layers.0.weight', 'output_blocks.8.0.out_layers.0.bias', 'output_blocks.8.0.out_layers.3.weight', 'output_blocks.8.0.out_layers.3.bias', 'output_blocks.8.0.skip_connection.weight', 'output_blocks.8.0.skip_connection.bias', 'output_blocks.8.1.norm.weight', 'output_blocks.8.1.norm.bias', 'output_blocks.8.1.qkv.weight', 'output_blocks.8.1.qkv.bias', 'output_blocks.8.1.proj_out.weight', 'output_blocks.8.1.proj_out.bias', 'output_blocks.8.2.conv.weight', 'output_blocks.8.2.conv.bias', 'output_blocks.9.0.in_layers.0.weight', 'output_blocks.9.0.in_layers.0.bias', 'output_blocks.9.0.in_layers.2.weight', 'output_blocks.9.0.in_layers.2.bias', 'output_blocks.9.0.emb_layers.1.weight', 'output_blocks.9.0.emb_layers.1.bias', 'output_blocks.9.0.out_layers.0.weight', 'output_blocks.9.0.out_layers.0.bias', 'output_blocks.9.0.out_layers.3.weight', 'output_blocks.9.0.out_layers.3.bias', 'output_blocks.9.0.skip_connection.weight', 'output_blocks.9.0.skip_connection.bias', 'output_blocks.10.0.in_layers.0.weight', 'output_blocks.10.0.in_layers.0.bias', 'output_blocks.10.0.in_layers.2.weight', 'output_blocks.10.0.in_layers.2.bias', 'output_blocks.10.0.emb_layers.1.weight', 'output_blocks.10.0.emb_layers.1.bias', 'output_blocks.10.0.out_layers.0.weight', 'output_blocks.10.0.out_layers.0.bias', 'output_blocks.10.0.out_layers.3.weight', 'output_blocks.10.0.out_layers.3.bias', 'output_blocks.10.0.skip_connection.weight', 'output_blocks.10.0.skip_connection.bias', 'output_blocks.11.0.in_layers.0.weight', 'output_blocks.11.0.in_layers.0.bias', 'output_blocks.11.0.in_layers.2.weight', 'output_blocks.11.0.in_layers.2.bias', 'output_blocks.11.0.emb_layers.1.weight', 'output_blocks.11.0.emb_layers.1.bias', 'output_blocks.11.0.out_layers.0.weight', 'output_blocks.11.0.out_layers.0.bias', 'output_blocks.11.0.out_layers.3.weight', 'output_blocks.11.0.out_layers.3.bias', 'output_blocks.11.0.skip_connection.weight', 'output_blocks.11.0.skip_connection.bias', 'output_blocks.11.1.conv.weight', 'output_blocks.11.1.conv.bias', 'output_blocks.12.0.in_layers.0.weight', 'output_blocks.12.0.in_layers.0.bias', 'output_blocks.12.0.in_layers.2.weight', 'output_blocks.12.0.in_layers.2.bias', 'output_blocks.12.0.emb_layers.1.weight', 'output_blocks.12.0.emb_layers.1.bias', 'output_blocks.12.0.out_layers.0.weight', 'output_blocks.12.0.out_layers.0.bias', 'output_blocks.12.0.out_layers.3.weight', 'output_blocks.12.0.out_layers.3.bias', 'output_blocks.12.0.skip_connection.weight', 'output_blocks.12.0.skip_connection.bias', 'output_blocks.13.0.in_layers.0.weight', 'output_blocks.13.0.in_layers.0.bias', 'output_blocks.13.0.in_layers.2.weight', 'output_blocks.13.0.in_layers.2.bias', 'output_blocks.13.0.emb_layers.1.weight', 'output_blocks.13.0.emb_layers.1.bias', 'output_blocks.13.0.out_layers.0.weight', 'output_blocks.13.0.out_layers.0.bias', 'output_blocks.13.0.out_layers.3.weight', 'output_blocks.13.0.out_layers.3.bias', 'output_blocks.13.0.skip_connection.weight', 'output_blocks.13.0.skip_connection.bias', 'output_blocks.14.0.in_layers.0.weight', 'output_blocks.14.0.in_layers.0.bias', 'output_blocks.14.0.in_layers.2.weight', 'output_blocks.14.0.in_layers.2.bias', 'output_blocks.14.0.emb_layers.1.weight', 'output_blocks.14.0.emb_layers.1.bias', 'output_blocks.14.0.out_layers.0.weight', 'output_blocks.14.0.out_layers.0.bias', 'output_blocks.14.0.out_layers.3.weight', 'output_blocks.14.0.out_layers.3.bias', 'output_blocks.14.0.skip_connection.weight', 'output_blocks.14.0.skip_connection.bias', 'out.0.weight', 'out.0.bias', 'out.2.weight', 'out.2.bias'])"
-      ]
-     },
-     "execution_count": 18,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "state0.keys()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "end\n"
-     ]
-    }
-   ],
-   "source": [
+    "state1 = model1['unet_state_dict']\n",
+    "state00 = model00['unet_state_dict']\n",
+    "state11 = model11['unet_state_dict']\n",
+    "state0.keys()\n",
     "for key in state0.keys():\n",
     "    # print(key)\n",
-    "    if not torch.equal(state1[key],state0[key]):\n",
+    "    if not torch.equal(state00[key],state11[key]):\n",
     "        print(key, \"different\")\n",
     "        # break\n",
     "    # else:\n",