models.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from numbers import Number\n",
    "import numpy as np\n",
    "import pdb\n",
    "import pickle\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "import torch.distributions as dists\n",
    "from torch_geometric.nn.conv import MessagePassing\n",
    "from torch_geometric.data import Data\n",
    "from typing import Union\n",
    "\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..', '..'))\n",
    "from concept_library.util import filter_kwargs, gather_broadcast, COLOR_LIST, Zip, init_args, to_np_array, get_filename_short, to_cpu\n",
    "from concept_library.util import extend_dims, record_data, transform_dict, to_cpu, to_device_recur, init_args, get_soft_Jaccard_distance, ddeepcopy as deepcopy\n",
    "from concept_library.util import get_activation, get_normalization, repeat_n, visualize_matrices, Shared_Param_Dict, MLP\n",
    "from concept_library.settings import REPR_DIM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Energy-based models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper functions\n",
    "def requires_grad(parameters, flag=True):\n",
    "    for p in parameters:\n",
    "        p.requires_grad = flag\n",
    "\n",
    "\n",
    "def id_to_tensor(ids, CONCEPTS, OPERATORS, requires_grad=False):\n",
    "    \"\"\"Concept string to repr tensor.\"\"\"\n",
    "    if not (isinstance(ids, list) or isinstance(ids, tuple)):\n",
    "        ids = [ids]\n",
    "    tensor = []\n",
    "    for id in ids:\n",
    "        if id in CONCEPTS:\n",
    "            tensor.append(CONCEPTS[id].get_node_repr())\n",
    "        elif id in OPERATORS:\n",
    "            tensor.append(OPERATORS[id].get_node_repr())\n",
    "        else:\n",
    "            raise\n",
    "    if requires_grad:\n",
    "        return torch.stack(tensor)\n",
    "    else:\n",
    "        return torch.stack(tensor).detach()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 Main models:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1.0 Load and get model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model_energy(args, device=\"cpu\"):\n",
    "    \"\"\"Get EBM model according to args.\"\"\"\n",
    "    if args.model_type == \"IGEBM\":\n",
    "        model = IGEBM(\n",
    "            in_channels=args.in_channels,\n",
    "            n_classes=args.n_classes,\n",
    "            channel_base=args.channel_base,\n",
    "            is_spec_norm=args.is_spec_norm,\n",
    "            aggr_mode=args.aggr_mode,\n",
    "        ).to(device)\n",
    "    elif args.model_type == \"CEBM\":\n",
    "        model = ConceptEBM(\n",
    "            mode=\"operator\" if args.is_two_branch else \"concept\",\n",
    "            in_channels=args.in_channels,\n",
    "            repr_dim=REPR_DIM,\n",
    "            w_type=args.w_type,\n",
    "            mask_mode=args.mask_mode,\n",
    "            channel_base=args.channel_base,\n",
    "            two_branch_mode=args.two_branch_mode,\n",
    "            is_spec_norm=args.is_spec_norm,\n",
    "            is_res=args.is_res,\n",
    "            c_repr_mode=args.c_repr_mode,\n",
    "            c_repr_first=args.c_repr_first,\n",
    "            z_mode=args.z_mode,\n",
    "            z_first=args.z_first,\n",
    "            z_dim=args.z_dim,\n",
    "            pos_embed_mode=args.pos_embed_mode,\n",
    "            aggr_mode=args.aggr_mode,\n",
    "            act_name=args.act_name,\n",
    "            normalization_type=args.normalization_type,\n",
    "            dropout=args.dropout,\n",
    "            self_attn_mode=args.self_attn_mode,\n",
    "            last_act_name=args.last_act_name,\n",
    "            n_avg_pool=args.n_avg_pool,\n",
    "        ).to(device)\n",
    "    elif args.model_type == \"CEBMLarge\":\n",
    "        model = ConceptEBMLarge(\n",
    "            mode=\"operator\" if args.is_two_branch else \"concept\",\n",
    "            in_channels=args.in_channels,\n",
    "            repr_dim=REPR_DIM,\n",
    "            w_type=args.w_type,\n",
    "            mask_mode=args.mask_mode,\n",
    "            channel_base=args.channel_base,\n",
    "            two_branch_mode=args.two_branch_mode,\n",
    "            is_spec_norm=args.is_spec_norm,\n",
    "            is_res=args.is_res,\n",
    "            c_repr_mode=args.c_repr_mode,\n",
    "            c_repr_first=args.c_repr_first,\n",
    "            z_mode=args.z_mode,\n",
    "            z_first=args.z_first,\n",
    "            z_dim=args.z_dim,\n",
    "            pos_embed_mode=args.pos_embed_mode,\n",
    "            aggr_mode=args.aggr_mode,\n",
    "            act_name=args.act_name,\n",
    "            normalization_type=args.normalization_type,\n",
    "            dropout=args.dropout,\n",
    "            self_attn_mode=args.self_attn_mode,\n",
    "            last_act_name=args.last_act_name,\n",
    "            is_multiscale=True,\n",
    "        ).to(device)\n",
    "    else:\n",
    "        raise Exception(\"model_type '{}' is not valid!\".format(args.model_type))\n",
    "    return model\n",
    "\n",
    "\n",
    "def load_model_energy(model_dict, device=\"cpu\"):\n",
    "    \"\"\"Load EBM model.\"\"\"\n",
    "    if model_dict is None:\n",
    "        return model_dict\n",
    "    model_type = model_dict[\"type\"]\n",
    "    if model_type == \"IGEBM\":\n",
    "        model = IGEBM(\n",
    "            in_channels=model_dict[\"in_channels\"],\n",
    "            n_classes=model_dict[\"n_classes\"],\n",
    "            channel_base=model_dict[\"channel_base\"] if \"channel_base\" in model_dict else 128,\n",
    "            is_spec_norm=model_dict[\"is_spec_norm\"] if \"is_spec_norm\" in model_dict else True,\n",
    "            aggr_mode=model_dict[\"aggr_mode\"] if \"aggr_mode\" in model_dict else \"sum\",\n",
    "        )\n",
    "    elif model_type == \"ConceptEBM\":\n",
    "        model = ConceptEBM(\n",
    "            mode=model_dict[\"mode\"] if \"mode\" in model_dict else \"concept\",\n",
    "            in_channels=model_dict[\"in_channels\"],\n",
    "            repr_dim=model_dict[\"repr_dim\"],\n",
    "            w_type=model_dict[\"w_type\"] if \"w_type\" in model_dict else \"image+mask\",\n",
    "            mask_mode=model_dict[\"mask_mode\"],\n",
    "            channel_base=model_dict[\"channel_base\"] if \"channel_base\" in model_dict else 128,\n",
    "            two_branch_mode=model_dict[\"two_branch_mode\"] if \"two_branch_mode\" in model_dict else \"concat\",\n",
    "            is_spec_norm=model_dict[\"is_spec_norm\"] if \"is_spec_norm\" in model_dict else \"True\",\n",
    "            is_res=model_dict[\"is_res\"] if \"is_res\" in model_dict else True,\n",
    "            c_repr_mode=model_dict[\"c_repr_mode\"] if \"c_repr_mode\" in model_dict else \"l1\",\n",
    "            c_repr_first=model_dict[\"c_repr_first\"] if \"c_repr_first\" in model_dict else 0,\n",
    "            c_repr_base=model_dict[\"c_repr_base\"] if \"c_repr_base\" in model_dict else 2,\n",
    "            z_mode=model_dict[\"z_mode\"] if \"z_mode\" in model_dict else \"None\",\n",
    "            z_first=model_dict[\"z_first\"] if \"z_first\" in model_dict else 2,\n",
    "            z_dim=model_dict[\"z_dim\"] if \"z_dim\" in model_dict else 4,\n",
    "            pos_embed_mode=model_dict[\"pos_embed_mode\"] if \"pos_embed_mode\" in model_dict else \"None\",\n",
    "            aggr_mode=model_dict[\"aggr_mode\"] if \"aggr_mode\" in model_dict else \"sum\",\n",
    "            act_name=model_dict[\"act_name\"] if \"act_name\" in model_dict else \"leakyrelu0.2\",\n",
    "            normalization_type=model_dict[\"normalization_type\"] if \"normalization_type\" in model_dict else \"None\",\n",
    "            dropout=model_dict[\"dropout\"] if \"dropout\" in model_dict else 0,\n",
    "            self_attn_mode=model_dict[\"self_attn_mode\"] if \"self_attn_mode\" in model_dict else \"None\",\n",
    "            last_act_name=model_dict[\"last_act_name\"] if \"last_act_name\" in model_dict else \"None\",\n",
    "            n_avg_pool=model_dict[\"n_avg_pool\"] if \"n_avg_pool\" in model_dict else 0,\n",
    "        )\n",
    "        if \"c_repr\" in model_dict:\n",
    "            model.set_c(model_dict[\"c_repr\"], model_dict[\"c_str\"])\n",
    "    elif model_type == \"ConceptEBMLarge\":\n",
    "        model = ConceptEBMLarge(\n",
    "            mode=model_dict[\"mode\"] if \"mode\" in model_dict else \"concept\",\n",
    "            in_channels=model_dict[\"in_channels\"],\n",
    "            repr_dim=model_dict[\"repr_dim\"],\n",
    "            w_type=model_dict[\"w_type\"] if \"w_type\" in model_dict else \"image+mask\",\n",
    "            mask_mode=model_dict[\"mask_mode\"],\n",
    "            channel_base=model_dict[\"channel_base\"] if \"channel_base\" in model_dict else 128,\n",
    "            two_branch_mode=model_dict[\"two_branch_mode\"] if \"two_branch_mode\" in model_dict else \"concat\",\n",
    "            is_spec_norm=model_dict[\"is_spec_norm\"] if \"is_spec_norm\" in model_dict else \"True\",\n",
    "            is_res=model_dict[\"is_res\"] if \"is_res\" in model_dict else False,\n",
    "            c_repr_mode=model_dict[\"c_repr_mode\"] if \"c_repr_mode\" in model_dict else \"l1\",\n",
    "            c_repr_first=model_dict[\"c_repr_first\"] if \"c_repr_first\" in model_dict else 0,\n",
    "            c_repr_base=model_dict[\"c_repr_base\"] if \"c_repr_base\" in model_dict else 2,\n",
    "            z_mode=model_dict[\"z_mode\"] if \"z_mode\" in model_dict else \"None\",\n",
    "            z_first=model_dict[\"z_first\"] if \"z_first\" in model_dict else 2,\n",
    "            z_dim=model_dict[\"z_dim\"] if \"z_dim\" in model_dict else 4,\n",
    "            pos_embed_mode=model_dict[\"pos_embed_mode\"] if \"pos_embed_mode\" in model_dict else \"None\",\n",
    "            aggr_mode=model_dict[\"aggr_mode\"] if \"aggr_mode\" in model_dict else \"sum\",\n",
    "            act_name=model_dict[\"act_name\"] if \"act_name\" in model_dict else \"leakyrelu0.2\",\n",
    "            normalization_type=model_dict[\"normalization_type\"] if \"normalization_type\" in model_dict else \"None\",\n",
    "            dropout=model_dict[\"dropout\"] if \"dropout\" in model_dict else 0,\n",
    "            self_attn_mode=model_dict[\"self_attn_mode\"] if \"self_attn_mode\" in model_dict else \"None\",\n",
    "            last_act_name=model_dict[\"last_act_name\"] if \"last_act_name\" in model_dict else \"None\",\n",
    "            is_multiscale=model_dict[\"is_multiscale\"],\n",
    "        )\n",
    "        if \"c_repr\" in model_dict:\n",
    "            model.set_c(model_dict[\"c_repr\"], model_dict[\"c_str\"])\n",
    "    elif model_type == \"SumEBM\":\n",
    "        model = SumEBM(*[load_model_energy(model_dict_ele) for model_dict_ele in model_dict[\"models\"]])\n",
    "    elif model_type == \"GraphEBM\":\n",
    "        model = GraphEBM(\n",
    "            models={key: load_model_energy(model_dict_ele) for key, model_dict_ele in model_dict[\"models\"].items()},\n",
    "            assign_dict=model_dict[\"assign_dict\"],\n",
    "            mask_arity=model_dict[\"mask_arity\"],\n",
    "        )\n",
    "    elif model_type == \"GNN_energy\":\n",
    "        model = GNN_energy(\n",
    "            mode=model_dict[\"mode\"],\n",
    "            is_zgnn_node=model_dict[\"is_zgnn_node\"],\n",
    "            edge_attr_size=model_dict[\"edge_attr_size\"],\n",
    "            aggr_mode=model_dict[\"aggr_mode\"],\n",
    "            n_GN_layers=model_dict[\"n_GN_layers\"],\n",
    "            n_neurons=model_dict[\"n_neurons\"],\n",
    "            GNN_output_size=model_dict[\"GNN_output_size\"],\n",
    "            mlp_n_layers=model_dict[\"mlp_n_layers\"],\n",
    "            gnn_normalization_type=model_dict[\"gnn_normalization_type\"],\n",
    "            activation=model_dict[\"activation\"],\n",
    "            recurrent=model_dict[\"recurrent\"],\n",
    "            cnn_output_size=model_dict[\"cnn_output_size\"],\n",
    "            cnn_is_spec_norm=model_dict[\"cnn_is_spec_norm\"],\n",
    "            cnn_normalization_type=model_dict[\"cnn_normalization_type\"],\n",
    "            cnn_channel_base=model_dict[\"cnn_channel_base\"],\n",
    "            cnn_aggr_mode=model_dict[\"cnn_aggr_mode\"],\n",
    "            c_repr_dim=model_dict[\"c_repr_dim\"],\n",
    "            z_dim=model_dict[\"z_dim\"],\n",
    "            zgnn_dim=model_dict[\"zgnn_dim\"],\n",
    "            distance_loss_type=model_dict[\"distance_loss_type\"],\n",
    "            pooling_type=model_dict[\"pooling_type\"],\n",
    "            pooling_dim=model_dict[\"pooling_dim\"],\n",
    "            is_x=model_dict[\"is_x\"] if \"is_x\" in model_dict else False,\n",
    "        )\n",
    "    else:\n",
    "        raise Exception(\"model_type '{}' is not valid!\".format(model_type))\n",
    "    if model_type not in [\"SumEBM\", \"GraphEBM\"]:\n",
    "        model.load_state_dict(model_dict[\"state_dict\"])\n",
    "    model.to(device)\n",
    "    return model\n",
    "\n",
    "\n",
    "def to_ebm_models(ebm_dict, device=\"cpu\"):\n",
    "    \"\"\"Transform every element into model from model_dict.\"\"\"\n",
    "    if len(ebm_dict) > 0:\n",
    "        if ebm_dict.__class__.__name__ == \"Shared_Param_Dict\":\n",
    "            return ebm_dict\n",
    "        elif \"type\" in ebm_dict and ebm_dict[\"type\"] == \"Shared_Param_Dict\":\n",
    "            ebm_dict = Shared_Param_Dict(\n",
    "                concept_model=load_model_energy(ebm_dict[\"concept_model_dict\"], device=device),\n",
    "                relation_model=load_model_energy(ebm_dict[\"relation_model_dict\"], device=device),\n",
    "                concept_repr_dict=ebm_dict[\"concept_repr_dict\"],\n",
    "                relation_repr_dict=ebm_dict[\"relation_repr_dict\"],\n",
    "            )\n",
    "        else:\n",
    "            first_item = ebm_dict[next(iter(ebm_dict))]\n",
    "            if isinstance(first_item, dict):\n",
    "                # Tranform from model_dict to actual model:\n",
    "                ebm_dict = {key: load_model_energy(model_dict, device=device) \n",
    "                            for key, model_dict in ebm_dict.items()}\n",
    "    return ebm_dict\n",
    "\n",
    "\n",
    "def load_best_model(\n",
    "    data_record,\n",
    "    keys=[\"mask|c_repr\", \"mask|c\", \"c_repr|mask\", \"c_repr|c\"],\n",
    "    return_id=False,\n",
    "    load_epoch=\"best\",\n",
    "):\n",
    "    \"\"\"Load best model according to the given acc keys.\"\"\"\n",
    "    args = init_args(update_default_hyperparam(data_record[\"args\"]))\n",
    "    acc_mean = np.array([data_record[\"acc\"][\"acc:{}:val\".format(key) if \"acc:{}:val\".format(key) in data_record[\"acc\"] else \"iou:{}:val\".format(key)] for key in keys]).mean(0)\n",
    "    acc_dict = {epoch: acc for epoch, acc in zip(data_record[\"acc\"][\"epoch:val\"], acc_mean) if epoch % args.save_interval == 0}\n",
    "    if load_epoch == \"best\":\n",
    "        acc_dict_argmax = np.argmax(list(acc_dict.values()))\n",
    "        acc_argmax_epoch = list(acc_dict.keys())[acc_dict_argmax]\n",
    "        best_model_id = data_record[\"save_epoch\"].index(acc_argmax_epoch)\n",
    "        best_model = load_model_energy(data_record[\"model_dict\"][best_model_id])\n",
    "        acc_chosen = np.max(list(acc_dict.values()))\n",
    "        assert acc_chosen == acc_dict[acc_argmax_epoch]\n",
    "        print(\"Loaded best model for {} at epoch {}. Best mean acc: {:.6f}\".format(args.dataset, acc_argmax_epoch, acc_chosen))\n",
    "    elif load_epoch == \"last\":\n",
    "        acc_argmax_epoch = data_record[\"save_epoch\"][-1]\n",
    "        best_model = load_model_energy(data_record[\"model_dict\"][-1])\n",
    "        acc_chosen = acc_dict[acc_argmax_epoch]\n",
    "        print(\"Loaded model for {} at last epoch {} with mean acc: {:.6f}. Best mean acc: {:.6f}\".format(args.dataset, acc_argmax_epoch, acc_chosen, np.max(list(acc_dict.values()))))\n",
    "    else:\n",
    "        acc_argmax_epoch = int(load_epoch)\n",
    "        best_model_id = data_record[\"save_epoch\"].index(acc_argmax_epoch)\n",
    "        best_model = load_model_energy(data_record[\"model_dict\"][best_model_id])\n",
    "        acc_chosen = acc_dict[acc_argmax_epoch]\n",
    "        print(\"Loaded best model for {} at epoch {} with mean acc: {:.6f}. Best mean acc: {:.6f}\".format(args.dataset, acc_argmax_epoch, acc_chosen, np.max(list(acc_dict.values()))))\n",
    "    if return_id:\n",
    "        return best_model, best_model_id\n",
    "    else:\n",
    "        return best_model\n",
    "\n",
    "\n",
    "def load_best_model_from_file(dirname, filename, keys=[\"mask|c_repr\", \"mask|c\", \"c_repr|mask\", \"c_repr|c\"], device=\"cpu\"):\n",
    "    \"\"\"Load best model according to the given acc keys from file.\"\"\"\n",
    "    data_record = pickle.load(open(dirname + filename, \"rb\"))\n",
    "    model = load_best_model(data_record, keys=keys).to(device)\n",
    "    model.set_repr_dict(filter_kwargs(data_record[\"concept_embeddings\"][-1], data_record[\"args\"][\"concept_collection\"]))\n",
    "    return model\n",
    "\n",
    "\n",
    "def load_model_atom(model_atom_str, model_type=\"CEBM\", device=\"cpu\"):\n",
    "    \"\"\"Load a dictionary of saved best models for atomic concepts and relations.\"\"\"\n",
    "    models = {}\n",
    "    for model_str in model_atom_str.split(\"^\"):\n",
    "        model_set = set(model_str.split(\"+\"))\n",
    "        model_mode = \"Load\"\n",
    "        if model_type == \"CEBM\":\n",
    "            if model_set.issubset({\"Line\", \"Lshape\", \"Rect\"}):\n",
    "                dirname = \"/dfs/user/tailin/.results/ebm_5-9.1/\"\n",
    "                filename = \"c-Line+Lshape+Rect_cz_8_model_CEBM_alpha_1_las_0.1_size_20.0_sams_60_e_r-rmb_et_mask_pl_False_nco_0.2_mask_concat_tbm_concat_cm_c2_cf_2_p_0.2_id_correct-detach_Hash_k2MkhWud_turing4.p\"\n",
    "            elif model_set.issubset({\"Vertical\", \"Parallel\"}):\n",
    "                dirname = \"/dfs/user/tailin/.results/ebm_5-9.1/\"\n",
    "                filename = \"c-Parallel+Vertical_cz_8_model_CEBM_alpha_1_las_0.1_size_20.0_sams_60_e_r-rmb_et_mask_pl_False_nco_0.2_mask_concat_tbm_concat_cm_c2_cf_2_p_0.2_id_1_Hash_tdH7khnR_turing2.p\"\n",
    "            elif model_set.issubset({\"SameAll\", \"SameShape\", \"SameColor\", \"SameRow\", \"SameCol\", \"IsInside\", \"IsTouch\"}):\n",
    "                dirname = \"/dfs/user/tailin/.results/ebm_5-9.1/\"\n",
    "                filename = \"cz_8_model_CEBM_alpha_1_las_0.1_size_20.0_sams_60_e_r-rmb_et_mask_pl_False_nco_0.2_mask_concat_tbm_concat_cm_c2_cf_2_p_0.2_id_correct-detach_Hash_tKR8Qvme_turing4.p\"\n",
    "            elif model_set.issubset({\"Image\"}):\n",
    "                dirname = \"/dfs/user/tailin/.results/ebm_5-9.1/\"\n",
    "                filename = \"c-Image_cz_8_model_CEBM_alpha_1_las_0.1_size_20.0_sams_60_e_r-rmb_et_mask_pl_False_nco_0.2_mask_concat_tbm_concat_cm_c2_cf_2_p_0.2_id_correct-detach_Hash_TkXJwZhk_turing4.p\"\n",
    "            elif model_set.issubset({\"RotateA\", \"RotateB\", \"RotateC\"}):\n",
    "                dirname = \"/dfs/user/tailin/.results/ebm_5-9.1/\"\n",
    "                filename = \"c-RotateA+RotateB+RotateC(Lshape)_cz_8_model_CEBM_alpha_1_las_0.1_size_20.0_sams_60_e_r-rmb_et_mask_pl_False_nco_0.2_mask_concat_tbm_concat_cm_c2_cf_2_p_0.2_id_1_Hash_f9Njs8Za_turing4.p\"\n",
    "            elif model_set.issubset({\"AdaptRe\"}):\n",
    "                model_mode = \"AdaptRe\"\n",
    "            else:\n",
    "                raise Exception(\"No saved model for {}.\".format(model_set))\n",
    "        else:\n",
    "            raise Exception(\"model_type {} is not valid!\".format(model_type))\n",
    "        if model_mode == \"Load\":\n",
    "            models[model_str] = load_best_model_from_file(dirname, filename, device=device)\n",
    "        elif model_mode == \"AdaptRe\":\n",
    "            \"\"\"Adaptive relations.\"\"\"\n",
    "            args_copy = deepcopy(args)\n",
    "            args_copy.is_two_branch = True\n",
    "            args_copy.model_type = \"CEBM\"\n",
    "            models[model_str] = get_model_energy(args_copy)\n",
    "        else:\n",
    "            raise\n",
    "    return models\n",
    "\n",
    "\n",
    "def update_default_hyperparam(Dict):\n",
    "    \"\"\"Default hyperparameters for previous experiments, after adding these new options.\"\"\"\n",
    "    default_param = {\n",
    "        \"is_two_branch\": False,\n",
    "        \"two_branch_mode\": \"concat\",\n",
    "        \"rainbow_prob\": 0,\n",
    "        \"max_n_distractors\": 0,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"n_operators\" : 1,\n",
    "        \"color_avail\" : \"-1\",\n",
    "        \"transforms\": \"None\",\n",
    "        \"transforms_pos\": \"None\",\n",
    "        # Training:\n",
    "        \"ebm_target_mode\": \"None\",\n",
    "        \"ebm_target\": \"mask\",\n",
    "        \"emp_target_mode\": \"all\",\n",
    "        \"is_pos_repr_learnable\": False,\n",
    "        \"p_buffer\": 0.95,\n",
    "        \"lambd_start\": -1,\n",
    "        \"lambd\": 0.005,\n",
    "        \"neg_mode\": \"None\",\n",
    "        \"neg_mode_coef\": 0.,\n",
    "        \"early_stopping_patience\": -1,\n",
    "        \"step_size_start\": -1,\n",
    "        \"step_size_img\": -1,\n",
    "        \"step_size_repr\": -1,\n",
    "        \"step_size_z\": 2,\n",
    "        \"step_size_zgnn\": 2,\n",
    "        \"step_size_wtarget\": -1,\n",
    "        \"is_spec_norm\": \"True\",\n",
    "        \"is_res\": True,\n",
    "        \"c_repr_mode\": \"l1\",\n",
    "        \"c_repr_first\": 0,\n",
    "        \"c_repr_base\": 2,\n",
    "        \"aggr_mode\": \"sum\",\n",
    "        \"act_name\": \"leakyrelu0.2\",\n",
    "        \"normalization_type\": \"None\",\n",
    "        \"dropout\": 0,\n",
    "        \"self_attn_mode\": \"None\",\n",
    "        \"last_act_name\": \"None\",\n",
    "        \"n_avg_pool\": 0,\n",
    "        \"kl_all_step\": False,\n",
    "        \"kl_coef\": 0.,\n",
    "        \"entropy_coef_img\": 0.,\n",
    "        \"entropy_coef_mask\": 0.,\n",
    "        \"entropy_coef_repr\": 0.,\n",
    "        \"epsilon_ent\": 1e-5,\n",
    "        \"pos_consistency_coef\": 0.,\n",
    "        \"neg_consistency_coef\": 0.,\n",
    "        \"emp_consistency_coef\": 0.,\n",
    "        # SGLD:\n",
    "        \"SGLD_is_anneal\": False,\n",
    "        \"SGLD_anneal_power\": 2.0,\n",
    "        \"SGLD_is_penalize_lower\": \"True\",\n",
    "        \"SGLD_mutual_exclusive_coef\": 0,\n",
    "        \"SGLD_fine_mutual_exclusive_coef\": 0,\n",
    "        \"SGLD_object_exceed_coef\": 0,\n",
    "        \"SGLD_pixel_entropy_coef\": 0,\n",
    "        \"SGLD_mask_entropy_coef\": 0,\n",
    "        \"SGLD_pixel_gm_coef\": 0,\n",
    "        # selector discovery:\n",
    "        \"SGLD_iou_batch_consistency_coef\": 0,\n",
    "        \"SGLD_iou_concept_repel_coef\": 0,\n",
    "        \"SGLD_iou_relation_repel_coef\": 0,\n",
    "        \"SGLD_iou_relation_overlap_coef\": 0,\n",
    "        \"SGLD_iou_attract_coef\": 0,\n",
    "        # Other settings:\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"train_mode\": \"cd\",\n",
    "        \"energy_mode\": \"standard\",\n",
    "        \"supervised_loss_type\": \"mse\",\n",
    "        \"target_loss_type\": \"mse\",\n",
    "        \"cumu_mode\": \"harmonic\",\n",
    "        \"channel_coef\": 1,\n",
    "        \"empty_coef\": 0.11,\n",
    "        \"obj_coef\": 0,\n",
    "        \"mutual_exclusive_coef\": 0,\n",
    "        \"pixel_entropy_coef\": 0,\n",
    "        \"pixel_gm_coef\": 0,\n",
    "        \"iou_batch_consistency_coef\": 0,\n",
    "        \"iou_concept_repel_coef\": 0,\n",
    "        \"iou_relation_repel_coef\": 0,\n",
    "        \"iou_relation_overlap_coef\": 0,\n",
    "        \"iou_attract_coef\": 0,\n",
    "        \"iou_target_matching_coef\": 0,\n",
    "        \"z_mode\": \"None\",\n",
    "        \"z_first\": 2,\n",
    "        \"z_dim\": 4,\n",
    "        \"pos_embed_mode\": \"None\",\n",
    "        \"image_value_range\": \"0,1\",\n",
    "        \"w_init_type\": \"random\",\n",
    "        \"indiv_sample\": -1,\n",
    "        \"n_tasks\": 128,\n",
    "        \"is_concat_minibatch\": False,\n",
    "        \"to_RGB\": False,\n",
    "        \"rescaled_size\": \"None\",\n",
    "        \"rescale_mode\": \"nearest\",\n",
    "        \"upsample\": -1,\n",
    "        \"relation_merge_mode\": \"None\",\n",
    "        \"is_relation_z\": True,\n",
    "        \"connected_coef\": 0,\n",
    "        \"connected_num_samples\": 2,\n",
    "        # Specific for EBM + GNN:\n",
    "        \"is_selector_gnn\": False,\n",
    "        \"is_zgnn_node\": False,\n",
    "        \"is_cross_validation\": False,\n",
    "        \"load_pretrained_concepts\": \"None\",\n",
    "        \"n_GN_layers\": 2,\n",
    "        \"gnn_normalization_type\": \"None\",\n",
    "        \"gnn_pooling_dim\": 16,\n",
    "        \"edge_attr_size\": 8,\n",
    "        \"cnn_output_size\": 32,\n",
    "        \"cnn_is_spec_norm\": \"True\",\n",
    "        \"train_coef\": 1,\n",
    "        \"test_coef\": 1,\n",
    "        \"lr_pretrained_concepts\": 0,\n",
    "        \"parallel_mode\": \"None\",\n",
    "        \"is_rewrite\": False,\n",
    "    }\n",
    "    for key, item in default_param.items():\n",
    "        if key not in Dict:\n",
    "            Dict[key] = item\n",
    "    return Dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1.1 IGEBM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class IGEBM(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels=3,\n",
    "        n_classes=None,\n",
    "        channel_base=128,\n",
    "        is_spec_norm=True,\n",
    "        aggr_mode=\"sum\",\n",
    "    ):\n",
    "        super(IGEBM, self).__init__()\n",
    "        self.in_channels = in_channels\n",
    "        self.n_classes = n_classes\n",
    "        self.channel_base = channel_base\n",
    "        self.aggr_mode = aggr_mode\n",
    "        self.is_spec_norm = is_spec_norm\n",
    "        if is_spec_norm:\n",
    "            self.conv1 = spectral_norm(nn.Conv2d(in_channels, channel_base, 3, padding=1), std=1)\n",
    "        else:\n",
    "            self.conv1 = nn.Conv2d(in_channels, channel_base, 3, padding=1)\n",
    "\n",
    "        self.blocks = nn.ModuleList(\n",
    "            [\n",
    "                ResBlock(channel_base, channel_base, n_classes, downsample=True, is_spec_norm=is_spec_norm),\n",
    "                ResBlock(channel_base, channel_base, n_classes, is_spec_norm=is_spec_norm),\n",
    "                ResBlock(channel_base, channel_base*2, n_classes, downsample=True, is_spec_norm=is_spec_norm),\n",
    "                ResBlock(channel_base*2, channel_base*2, n_classes, is_spec_norm=is_spec_norm),\n",
    "                ResBlock(channel_base*2, channel_base*2, n_classes, downsample=True, is_spec_norm=is_spec_norm),\n",
    "                ResBlock(channel_base*2, channel_base*2, n_classes, is_spec_norm=is_spec_norm),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.linear = nn.Linear(channel_base*2, 1)\n",
    "\n",
    "    def forward(self, input, class_id=None):\n",
    "        out = self.conv1(input)\n",
    "\n",
    "        out = F.leaky_relu(out, negative_slope=0.2)\n",
    "\n",
    "        for block in self.blocks:\n",
    "            out = block(out, class_id)\n",
    "\n",
    "        out = F.relu(out)\n",
    "        if self.aggr_mode == \"sum\":\n",
    "            out = out.view(out.shape[0], out.shape[1], -1).sum(2)\n",
    "        elif self.aggr_mode == \"max\":\n",
    "            out = out.view(out.shape[0], out.shape[1], -1).max(2)[0]\n",
    "        elif self.aggr_mode == \"mean\":\n",
    "            out = out.view(out.shape[0], out.shape[1], -1).mean(2)\n",
    "        else:\n",
    "            raise\n",
    "        out = self.linear(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"IGEBM\"}\n",
    "        model_dict[\"in_channels\"] = self.in_channels\n",
    "        model_dict[\"n_classes\"] = self.n_classes\n",
    "        model_dict[\"channel_base\"] = self.channel_base\n",
    "        model_dict[\"is_spec_norm\"] = self.is_spec_norm\n",
    "        model_dict[\"aggr_mode\"] = self.aggr_mode\n",
    "        model_dict[\"state_dict\"] = to_cpu(self.state_dict())\n",
    "        return model_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1.2 ConceptEBM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConceptEBM(nn.Module):\n",
    "    \"\"\"\n",
    "    An EBM designed specifically to find Concepts, e.g. Rect, Line, RectSolid or\n",
    "        Relations, e.g. SameShape, SameColor.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(\n",
    "        self,\n",
    "        mode=\"concept\",\n",
    "        in_channels=10,\n",
    "        repr_dim=4,\n",
    "        w_type=\"image+mask\",\n",
    "        mask_mode=\"concat\",\n",
    "        channel_base=128,\n",
    "        two_branch_mode=\"concat\",\n",
    "        is_spec_norm=True,\n",
    "        is_res=True,\n",
    "        c_repr_mode=\"l1\",\n",
    "        c_repr_first=0,\n",
    "        c_repr_base=2,\n",
    "        z_mode=\"None\",\n",
    "        z_first=2,\n",
    "        z_dim=4,\n",
    "        pos_embed_mode=\"None\",\n",
    "        aggr_mode=\"sum\",\n",
    "        img_dims=2,\n",
    "        act_name=\"leakyrelu0.2\",\n",
    "        normalization_type=\"None\",\n",
    "        dropout=0,\n",
    "        self_attn_mode=\"None\",\n",
    "        last_act_name=\"None\",\n",
    "        n_avg_pool=0,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Initialize the EBM.\n",
    "\n",
    "        Args:\n",
    "            w_type: type of the first two arities of input.\n",
    "                choose from \"image\", \"mask\", \"image+mask\", \"obj\", \"image+obj\"\n",
    "        \"\"\"\n",
    "        super(ConceptEBM, self).__init__()\n",
    "        self.mode = mode\n",
    "        self.in_channels = in_channels\n",
    "        self.repr_dim = repr_dim\n",
    "        self.w_type = w_type\n",
    "        self.w_dim = 1 if \"mask\" in self.w_type else in_channels\n",
    "        self.mask_mode = mask_mode\n",
    "        self.channel_base = channel_base\n",
    "        self.two_branch_mode = two_branch_mode\n",
    "\n",
    "        if self.mode in [\"concept\"]:\n",
    "            self.mask_arity = 1\n",
    "        elif self.mode in [\"operator\"]:\n",
    "            self.mask_arity = 2\n",
    "        else:\n",
    "            raise\n",
    "        self.is_spec_norm = is_spec_norm\n",
    "        self.is_res = is_res\n",
    "        self.c_repr_mode = c_repr_mode\n",
    "        self.c_repr_first = c_repr_first\n",
    "        self.c_repr_base = c_repr_base\n",
    "        self.z_mode = z_mode\n",
    "        self.z_first = z_first\n",
    "        self.z_dim = z_dim\n",
    "        self.pos_embed_mode = pos_embed_mode\n",
    "        self.aggr_mode = aggr_mode\n",
    "        self.img_dims = img_dims\n",
    "        self.act_name = act_name\n",
    "        self.act = get_activation(act_name)\n",
    "        self.normalization_type = normalization_type\n",
    "        self.dropout = dropout\n",
    "        assert self_attn_mode in [\"None\", \"pixel\"]\n",
    "        self.self_attn_mode = self_attn_mode\n",
    "        self.last_act_name = last_act_name\n",
    "        self.n_avg_pool = n_avg_pool\n",
    "        if img_dims == 2:\n",
    "            kernel_size = 3\n",
    "            padding = 1\n",
    "        elif img_dims == 1:\n",
    "            kernel_size = (3, 1)\n",
    "            padding = (1, 0)\n",
    "        else:\n",
    "            raise\n",
    "        if is_spec_norm in [True, \"True\"]:\n",
    "            self.conv1 = spectral_norm(nn.Conv2d(in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base if self.mode==\"concept\" else channel_base//2, kernel_size, padding=padding), std=1)\n",
    "        elif is_spec_norm in [False, \"False\"]:\n",
    "            self.conv1 = nn.Conv2d(in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base if self.mode==\"concept\" else channel_base//2, kernel_size, padding=padding)\n",
    "        elif is_spec_norm == \"ws\":\n",
    "            self.conv1 = WSConv2d(in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base if self.mode==\"concept\" else channel_base//2, kernel_size, padding=padding)\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        if self.mode in [\"concept\"] or self.two_branch_mode == \"concat\":\n",
    "            self.blocks = nn.ModuleList(\n",
    "                [\n",
    "                    CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    CResBlock(channel_base, channel_base, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    CResBlock(channel_base, channel_base*2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=2 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=2 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=3 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=3 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                ]\n",
    "            )\n",
    "            if self_attn_mode != \"None\":\n",
    "                assert self_attn_mode in [\"pixel\"]\n",
    "                self.attn = Self_Attn(in_dim=channel_base, act_name=act_name)\n",
    "        elif self.two_branch_mode.startswith(\"imbal\"):\n",
    "            n_indi_layers = int(self.two_branch_mode.split(\"-\")[1])\n",
    "            if self_attn_mode != \"None\":\n",
    "                assert self_attn_mode in [\"pixel\"]\n",
    "                self.attn_0 = Self_Attn(in_dim=channel_base//2, act_name=act_name)\n",
    "                self.attn_1 = Self_Attn(in_dim=channel_base//2, act_name=act_name)\n",
    "            if n_indi_layers == 1:\n",
    "                self.blocks_0 = nn.ModuleList(\n",
    "                    [\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    ]\n",
    "                )\n",
    "                self.blocks_1 = nn.ModuleList(\n",
    "                    [\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    ]\n",
    "                )\n",
    "                self.blocks = nn.ModuleList(\n",
    "                    [\n",
    "                        CResBlock(channel_base, channel_base*2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=2 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=2 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=3 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=3 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    ]\n",
    "                )\n",
    "            elif n_indi_layers == 2:\n",
    "                self.blocks_0 = nn.ModuleList(\n",
    "                    [\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base//2, channel_base, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base, channel_base, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    ]\n",
    "                )\n",
    "                self.blocks_1 = nn.ModuleList(\n",
    "                    [\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base//2, channel_base, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base, channel_base, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    ]\n",
    "                )\n",
    "                self.blocks = nn.ModuleList(\n",
    "                    [\n",
    "                        CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=2 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=2 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=3 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=3 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    ]\n",
    "                )\n",
    "            elif n_indi_layers == 3:\n",
    "                self.blocks_0 = nn.ModuleList(\n",
    "                    [\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base//2, channel_base, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base, channel_base, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=2 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=2 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    ]\n",
    "                )\n",
    "                self.blocks_1 = nn.ModuleList(\n",
    "                    [\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base//2, channel_base, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base, channel_base, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                        CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=2 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=2 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    ]\n",
    "                )\n",
    "                self.blocks = nn.ModuleList(\n",
    "                    [\n",
    "                        CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, is_spec_norm=is_spec_norm, is_res=is_res, c_repr_mode=c_repr_mode if c_repr_first<=3 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=3 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout),\n",
    "                    ]\n",
    "                )\n",
    "        else:\n",
    "            raise Exception(\"two_branch_mode '{}' is invalid!\".format(self.two_branch_mode))\n",
    "\n",
    "        self.linear = nn.Linear(channel_base*2, 1)\n",
    "\n",
    "    def forward(self, input, mask, c_repr=None, z=None, **kwargs):\n",
    "        \"\"\"\n",
    "        Returns the energy of a given mask on top of a given input.\n",
    "\n",
    "        Args:\n",
    "            input: input of shape [B, C, H, W]\n",
    "            mask: mask of shape [B, c, H, W]\n",
    "            c_repr: shape [B, repr_dim], explicit concept representation to use for the mask. If not set, uses self.c_repr.\n",
    "            z: a 1-tuple, element has shape [B, z_dim], latent representation depending on input and mask.\n",
    "\n",
    "        Returns:\n",
    "            out: energy, with shape [B, 1]\n",
    "        \"\"\"\n",
    "        length = len(mask[0])\n",
    "        if c_repr is None:\n",
    "            c_repr = self.c_repr.expand(length, self.c_repr.shape[1])\n",
    "        else:\n",
    "            if c_repr.shape[0] == 1:\n",
    "                c_repr = c_repr.expand(length, c_repr.shape[1])\n",
    "        c_repr_first_dict = {0: 0, 1: 2, 2: 4, 3: 5}\n",
    "        z_first_dict = {0: 0, 1: 2, 2: 4, 3: 5}\n",
    "        if self.mode == \"concept\":\n",
    "            assert len(mask) == 1\n",
    "            if self.mask_mode == \"concat\":\n",
    "                input_aug = torch.cat([input, mask[0]], 1)  # input: [B, C, H, W], mask: [B, c, H, W]\n",
    "            elif self.mask_mode == \"mul\":\n",
    "                input_aug = input * mask[0]\n",
    "            elif self.mask_mode == \"mulcat\":\n",
    "                input_aug = torch.cat([input*mask[0], mask[0]], 1)\n",
    "            elif self.mask_mode == \"fil\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_aug = torch.cat([torch.clamp(input[:,:1]+(1-mask[0]), 0, 1), input[:,1:]*mask[0]], 1)\n",
    "            elif self.mask_mode == \"filcat\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_fil = torch.cat([torch.clamp(input[:,:1]+(1-mask[0]), 0, 1), input[:,1:]*mask[0]], 1)\n",
    "                input_aug = torch.cat([input_fil, mask[0]], 1)\n",
    "            else:\n",
    "                raise\n",
    "            for _ in range(self.n_avg_pool):\n",
    "                input_aug = F.avg_pool2d(input_aug, 3, stride=2, padding=1)\n",
    "            out = self.conv1(input_aug)\n",
    "\n",
    "            out = self.act(out)\n",
    "\n",
    "            for i, block in enumerate(self.blocks):\n",
    "                if hasattr(self, \"self_attn_mode\") and self.self_attn_mode != \"None\" and i == 2:\n",
    "                    out, _ = self.attn(out)\n",
    "                out = block(\n",
    "                    out,\n",
    "                    c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=i else None,\n",
    "                    z=z if z_first_dict[self.z_first]<=i else None,\n",
    "                )\n",
    "\n",
    "            out = F.relu(out)\n",
    "            if self.aggr_mode == \"sum\":\n",
    "                out = out.view(out.shape[0], out.shape[1], -1).sum(2)\n",
    "            elif self.aggr_mode == \"max\":\n",
    "                out = out.view(out.shape[0], out.shape[1], -1).max(2)[0]\n",
    "            elif self.aggr_mode == \"mean\":\n",
    "                out = out.view(out.shape[0], out.shape[1], -1).mean(2)\n",
    "            else:\n",
    "                raise\n",
    "            out = self.linear(out)\n",
    "\n",
    "        elif self.mode == \"operator\":\n",
    "            # Combining input image and mask:\n",
    "            if not (isinstance(input, tuple) or isinstance(input, list)):\n",
    "                input = (input, input)\n",
    "            if self.mask_mode == \"concat\":\n",
    "                input_aug_0 = torch.cat([input[0], mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([input[1], mask[1]], 1)\n",
    "            elif self.mask_mode == \"mul\":\n",
    "                input_aug_0 = input[0] * mask[0]\n",
    "                input_aug_1 = input[1] * mask[1]\n",
    "            elif self.mask_mode == \"mulcat\":\n",
    "                input_aug_0 = torch.cat([input[0]*mask[0], mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([input[1]*mask[1], mask[1]], 1)\n",
    "            elif self.mask_mode == \"fil\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_aug_0 = torch.cat([torch.clamp(input[0][:,:1]+(1-mask[0]), 0, 1), input[0][:,1:]*mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([torch.clamp(input[1][:,:1]+(1-mask[1]), 0, 1), input[1][:,1:]*mask[1]], 1)\n",
    "            elif self.mask_mode == \"filcat\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_fil_0 = torch.cat([torch.clamp(input[0][:,:1]+(1-mask[0]), 0, 1), input[0][:,1:]*mask[0]], 1)\n",
    "                input_aug_0 = torch.cat([input_fil_0, mask[0]], 1)\n",
    "                input_fil_1 = torch.cat([torch.clamp(input[1][:,:1]+(1-mask[1]), 0, 1), input[1][:,1:]*mask[1]], 1)\n",
    "                input_aug_1 = torch.cat([input_fil_1, mask[1]], 1)\n",
    "            else:\n",
    "                raise\n",
    "\n",
    "            for _ in range(self.n_avg_pool):\n",
    "                input_aug_0 = F.avg_pool2d(input_aug_0, 3, stride=2, padding=1)\n",
    "                input_aug_1 = F.avg_pool2d(input_aug_1, 3, stride=2, padding=1)\n",
    "\n",
    "            out_0 = self.act(self.conv1(input_aug_0))\n",
    "            out_1 = self.act(self.conv1(input_aug_1))\n",
    "\n",
    "            if self.two_branch_mode == \"concat\":\n",
    "                out = torch.cat([out_0, out_1], 1)\n",
    "\n",
    "                for i, block in enumerate(self.blocks):\n",
    "                    if hasattr(self, \"self_attn_mode\") and self.self_attn_mode != \"None\" and i == 2:\n",
    "                        out, _ = self.attn(out)\n",
    "                    out = block(\n",
    "                        out,\n",
    "                        c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=i else None,\n",
    "                        z=z if z_first_dict[self.z_first]<=i else None,\n",
    "                    )\n",
    "\n",
    "                out = F.relu(out)  # [B, 128, 1, 1]\n",
    "                if self.aggr_mode == \"sum\":\n",
    "                    out = out.view(out.shape[0], out.shape[1], -1).sum(2)\n",
    "                elif self.aggr_mode == \"max\":\n",
    "                    out = out.view(out.shape[0], out.shape[1], -1).max(2)[0]\n",
    "                elif self.aggr_mode == \"mean\":\n",
    "                    out = out.view(out.shape[0], out.shape[1], -1).mean(2)\n",
    "                else:\n",
    "                    raise\n",
    "                out = self.linear(out)  # [B, 1]\n",
    "            elif self.two_branch_mode.startswith(\"imbal\"):\n",
    "                default_c_repr = torch.ones_like(c_repr).to(c_repr.device)\n",
    "                default_z = torch.ones_like(z).to(z.device)\n",
    "\n",
    "                for i, block_0 in enumerate(self.blocks_0):\n",
    "                    if hasattr(self, \"self_attn_mode\") and self.self_attn_mode != \"None\" and i == 2:\n",
    "                        out_0, _ = self.attn_0(out_0)\n",
    "                    out_0 = block_0(\n",
    "                        out_0,\n",
    "                        c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=i else None,\n",
    "                        z=z if z_first_dict[self.z_first]<=i else None,\n",
    "                    )\n",
    "                for i, block_1 in enumerate(self.blocks_1):\n",
    "                    if hasattr(self, \"self_attn_mode\") and self.self_attn_mode != \"None\" and i == 2:\n",
    "                        out_1, _ = self.attn_1(out_1)\n",
    "                    out_1 = block_1(\n",
    "                        out_1,\n",
    "                        c_repr=default_c_repr if c_repr_first_dict[self.c_repr_first]<=i else None,\n",
    "                        z=default_z if z_first_dict[self.z_first]<=i else None,\n",
    "                    )\n",
    "                out = torch.cat([out_0, out_1], 1)\n",
    "                for i, block in enumerate(self.blocks):\n",
    "                    out = block(\n",
    "                        out,\n",
    "                        c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=i+len(self.blocks_0) else None,\n",
    "                        z=z if z_first_dict[self.z_first]<=i+len(self.blocks_0) else None,\n",
    "                    )\n",
    "                out = F.relu(out)\n",
    "                if self.aggr_mode == \"sum\":\n",
    "                    out = out.view(out.shape[0], out.shape[1], -1).sum(2)\n",
    "                elif self.aggr_mode == \"max\":\n",
    "                    out = out.view(out.shape[0], out.shape[1], -1).max(2)[0]\n",
    "                elif self.aggr_mode == \"mean\":\n",
    "                    out = out.view(out.shape[0], out.shape[1], -1).mean(2)\n",
    "                else:\n",
    "                    raise\n",
    "\n",
    "                out = self.linear(out)\n",
    "\n",
    "        else:\n",
    "            raise Exception(\"self.mode {} is not supported!\".format(self.mode))\n",
    "\n",
    "        # Last activation:\n",
    "        if hasattr(self, \"last_act_name\") and self.last_act_name != \"None\":\n",
    "            if self.last_act_name in [\"square\", \"softplus\", \"exp\", \"sigmoid\"]:\n",
    "                out = get_activation(self.last_act_name)(out)\n",
    "            else:\n",
    "                raise\n",
    "        return out\n",
    "\n",
    "    def classify(self, input, mask, concept_collection, topk=-1, CONCEPTS=None, OPERATORS=None):\n",
    "        \"\"\"\n",
    "        Given the input and mask, classify the selected concept by picking the\n",
    "        lowest-energy concept from concept_collection.\n",
    "        \"\"\"\n",
    "        if isinstance(input, tuple) or isinstance(input, list):\n",
    "            length = len(input[0])\n",
    "            device = input[0].device\n",
    "        else:\n",
    "            length = len(input)\n",
    "            device = input.device\n",
    "        if topk == -1:\n",
    "            topk = len(concept_collection)\n",
    "        c_repr_energy = []\n",
    "        for j in range(len(concept_collection)):\n",
    "            c_repr = id_to_tensor([concept_collection[j]] * length, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device) # len 4\n",
    "            neg_energy = self(input, mask=mask, c_repr=c_repr)\n",
    "            c_repr_energy.append(neg_energy)\n",
    "        c_repr_energy = torch.cat(c_repr_energy, 1)\n",
    "        c_repr_argsort = c_repr_energy.argsort(1)\n",
    "        c_repr_pred_list = []\n",
    "        for i, argsort in enumerate(c_repr_argsort):\n",
    "            c_repr_pred = {}\n",
    "            for k in range(min(topk, len(concept_collection))):\n",
    "                id_k = c_repr_argsort[i][k]\n",
    "                c_repr_pred[concept_collection[id_k]] = c_repr_energy[i][id_k].item()\n",
    "            c_repr_pred_list.append(c_repr_pred)\n",
    "        return c_repr_pred_list\n",
    "\n",
    "    def ground(\n",
    "        self,\n",
    "        input,\n",
    "        args,\n",
    "        mask=None,\n",
    "        c_repr=None,\n",
    "        z=None,\n",
    "        ensemble_size=18,\n",
    "        topk=-1,\n",
    "        w_init_type=\"random\",\n",
    "        sample_step=150,\n",
    "        is_grad=False,\n",
    "        is_return_E=False,\n",
    "        isplot=True,\n",
    "        **kwargs\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Given an input image, find the best mask to match the given (optional)\n",
    "        concept representation.\n",
    "\n",
    "        Internally, uses a model ensemble for best results.\n",
    "        \"\"\"\n",
    "        \n",
    "        def init_neg_mask(input, init, ensemble_size):\n",
    "            \"\"\"Initialize negative mask\"\"\"\n",
    "            if isinstance(input, tuple):\n",
    "                assert len(input[0].shape) == len(input[1].shape) == 4\n",
    "                device = input[0].device\n",
    "                w_dim = 1 if \"mask\" in self.w_type else input[0].shape[1]\n",
    "                neg_mask = (torch.rand(input[0].shape[0]*ensemble_size, w_dim, *input[0].shape[2:]).to(device), torch.rand(input[1].shape[0]*ensemble_size, w_dim, *input[1].shape[2:]).to(device))\n",
    "                if init == \"input-mask\":\n",
    "                    assert input[0].shape[1] == 10\n",
    "                    input_l = repeat_n(input, n_repeats=ensemble_size)\n",
    "                    neg_mask = (neg_mask[0] * (input_l[0].argmax(1)[:, None] != 0), neg_mask[1] * (input_l[1].argmax(1)[:, None] != 0))\n",
    "                neg_mask[0].requires_grad = True\n",
    "                neg_mask[1].requires_grad = True\n",
    "            else:\n",
    "                assert len(input.shape) == 4\n",
    "                device = input.device\n",
    "                w_dim = 1 if \"mask\" in self.w_type else input.shape[1]\n",
    "                neg_mask = (torch.rand(input.shape[0]*ensemble_size, w_dim, *input.shape[2:]).to(device),)\n",
    "                if init == \"input-mask\":\n",
    "                    assert input.shape[1] == 10\n",
    "                    input_l = repeat_n(input, n_repeats=ensemble_size)\n",
    "                    neg_mask = (neg_mask[0] * (input_l.argmax(1)[:, None] != 0),)\n",
    "                neg_mask[0].requires_grad = True\n",
    "            return neg_mask\n",
    "\n",
    "        # Update args:\n",
    "        args = deepcopy(args)\n",
    "        for key, value in kwargs.items():\n",
    "            setattr(args, key, value)\n",
    "        args.sample_step = sample_step\n",
    "        if isinstance(input, tuple) or isinstance(input, list):\n",
    "            args.is_image_tuple = True\n",
    "            device = input[0].device\n",
    "        else:\n",
    "            args.is_image_tuple = False\n",
    "            device = input.device\n",
    "\n",
    "        # Perform SGLD:\n",
    "        if mask is None:\n",
    "            neg_mask = init_neg_mask(input, init=w_init_type, ensemble_size=ensemble_size)\n",
    "        else:\n",
    "            neg_mask = tuple([repeat_n(mask[0], n_repeats=ensemble_size)])\n",
    "\n",
    "        if self.z_mode != \"None\":\n",
    "            if z is None:\n",
    "                z = tuple([torch.rand(neg_mask[0].shape[0], self.z_dim, device=device)])\n",
    "            else:\n",
    "                z = tuple([repeat_n(z[0], n_repeats=ensemble_size)])\n",
    "\n",
    "        (img_ensemble, neg_mask_ensemble, z_ensemble, zgnn_ensemble, wtarget_ensemble), neg_out_list_ensemble, info_ensemble = neg_mask_sgd_ensemble(\n",
    "            self, input, neg_mask, c_repr, z=z, zgnn=None, wtarget=None, args=args,\n",
    "            ensemble_size=ensemble_size, is_grad=is_grad,\n",
    "            is_return_E=is_return_E,\n",
    "        )\n",
    "\n",
    "        neg_out_ensemble = neg_out_list_ensemble[-1]  # neg_out_ensemble: [ensemble_size, B]\n",
    "        # Sort the obtained results by energy for each example:\n",
    "        neg_out_ensemble = torch.FloatTensor(neg_out_ensemble).transpose(0,1)  # [B, ensemble_size]\n",
    "        neg_out_argsort = neg_out_ensemble.argsort(1)  # [B, ensemble_size]\n",
    "        batch_size = neg_out_argsort.shape[0]\n",
    "        neg_out_ensemble_sorted = torch.stack([neg_out_ensemble[i][neg_out_argsort[i]] for i in range(batch_size)])\n",
    "        if zgnn_ensemble is not None or wtarget_ensemble is not None:\n",
    "            neg_task_out_ensemble = neg_out_ensemble.reshape(*batch_shape, -1).mean(1)  # [B_task, ensemble_size]\n",
    "            neg_task_out_argsort = neg_task_out_ensemble.argsort(1)\n",
    "\n",
    "        if img_ensemble is not None:\n",
    "            if args.is_image_tuple:\n",
    "                img_ensemble = tuple(img_ensemble[k].transpose(0,1) for k in range(len(img)))  # Each element [B, ensemble_size, C, H, W]\n",
    "                img_ensemble_sorted = []\n",
    "                for k in range(len(img)):\n",
    "                    img_ensemble_sorted.append(torch.stack([img_ensemble[k][i][neg_out_argsort[i]] for i in range(batch_size)]))\n",
    "                img_ensemble_sorted = tuple(img_ensemble_sorted)  # each element: [B, ensemble_size, C, H, W] sorted along dim=1 according to neg_out\n",
    "            else:\n",
    "                img_ensemble = img_ensemble.transpose(0,1)  # [B, ensemble_size, C, H, W]\n",
    "                img_ensemble_sorted = torch.stack([img_ensemble[i][neg_out_argsort[i]] for i in range(batch_size)])\n",
    "        else:\n",
    "            img_ensemble_sorted = None\n",
    "\n",
    "        if neg_mask_ensemble is not None:\n",
    "            neg_mask_ensemble = tuple(neg_mask_ensemble[k].transpose(0,1) for k in range(self.mask_arity))  # Each element [B, ensemble_size, C, H, W]\n",
    "            neg_mask_ensemble_sorted = []\n",
    "            for k in range(self.mask_arity):\n",
    "                neg_mask_ensemble_sorted.append(torch.stack([neg_mask_ensemble[k][i][neg_out_argsort[i]] for i in range(len(neg_mask_ensemble[0]))]))\n",
    "            neg_mask_ensemble_sorted = tuple(neg_mask_ensemble_sorted)  # each element: [B, ensemble_size, C, H, W] sorted along dim=1 according to neg_out\n",
    "        else:\n",
    "            neg_mask_ensemble_sorted = None\n",
    "\n",
    "        if z_ensemble is not None:\n",
    "            z_ensemble = tuple(z_ensemble[k].transpose(0,1) for k in range(len(z_ensemble)))  # Each element [B, ensemble_size, Z]\n",
    "            z_ensemble_sorted = []\n",
    "            for k in range(len(z_ensemble)):\n",
    "                z_ensemble_sorted.append(torch.stack([z_ensemble[k][i][neg_out_argsort[i]] for i in range(batch_size)]))\n",
    "            z_ensemble_sorted = tuple(z_ensemble_sorted)  # each element: [B, ensemble_size, Z] sorted along dim=1 according to neg_out\n",
    "        else:\n",
    "            z_ensemble_sorted = None\n",
    "\n",
    "        if zgnn_ensemble is not None:\n",
    "            zgnn_ensemble = tuple(zgnn_ensemble[k].transpose(0,1) if zgnn_ensemble[k] is not None else None for k in range(len(zgnn_ensemble)))  # Each element [B, ensemble_size, Zgnn]\n",
    "            zgnn_ensemble_sorted = []\n",
    "            for k in range(len(zgnn_ensemble)):\n",
    "                if zgnn_ensemble[k] is not None:\n",
    "                    zgnn_ensemble_sorted.append(torch.stack([zgnn_ensemble[k][i][neg_task_out_argsort[i]][:topk_core] for i in range(batch_shape[0])]))\n",
    "                else:\n",
    "                    zgnn_ensemble_sorted.append(None)\n",
    "            zgnn_ensemble_sorted = tuple(zgnn_ensemble_sorted)  # each element: [B, ensemble_size, zgnn_dim] sorted along dim=1 according to neg_out\n",
    "        else:\n",
    "            zgnn_ensemble_sorted = None\n",
    "\n",
    "        if wtarget_ensemble is not None:\n",
    "            wtarget_ensemble = wtarget_ensemble.transpose(0,1)  # [B, ensemble_size, w_dim, H, W]\n",
    "            wtarget_ensemble_sorted = torch.stack([wtarget_ensemble[i][neg_out_argsort[i]][:topk_core] for i in range(batch_size)])\n",
    "        else:\n",
    "            wtarget_ensemble_sorted = None\n",
    "\n",
    "        if isplot >= 2:\n",
    "            # Plot SGLD learning curve:\n",
    "            print(\"SGLD learning curve:\")\n",
    "            plt.figure(figsize=(12,6))\n",
    "            for i in range(min(neg_out_list_ensemble.shape[-1], 6)):  # neg_out_list_ensemble: [sample_step, ensemble_size, B]\n",
    "#                 print(\"Example {}\".format(i))\n",
    "                for k in range(min(6, ensemble_size)):\n",
    "                    plt.plot(neg_out_list_ensemble[:, neg_out_argsort[i][k],i], c=COLOR_LIST[k], label=\"id_{}\".format(k), alpha=0.4)\n",
    "            plt.legend()\n",
    "            plt.show()\n",
    "        return (img_ensemble_sorted, neg_mask_ensemble_sorted, z_ensemble_sorted, zgnn_ensemble_sorted, wtarget_ensemble_sorted), neg_out_ensemble_sorted\n",
    "\n",
    "    def set_c(self, c_repr, c_str=None):\n",
    "        \"\"\"Set default c_repr.\"\"\"\n",
    "        assert len(c_repr.shape) == 2 and c_repr.shape[0] == 1\n",
    "        if not isinstance(c_repr, torch.Tensor):\n",
    "            c_repr = torch.FloatTensor(c_repr)\n",
    "        device = next(iter(self.parameters())).device\n",
    "        self.c_repr = c_repr.to(device)\n",
    "        self.c_str = c_str\n",
    "        return self\n",
    "\n",
    "    def set_repr_dict(self, c_repr_dict):\n",
    "        \"\"\"Set the dictionary of c_repr.\"\"\"\n",
    "        self.c_repr_dict = {}\n",
    "        device = next(iter(self.parameters())).device\n",
    "        for key, c_repr in c_repr_dict.items():\n",
    "            if not isinstance(c_repr, torch.Tensor):\n",
    "                c_repr = torch.FloatTensor(c_repr)\n",
    "            self.c_repr_dict[key] = c_repr.to(device)\n",
    "        return self\n",
    "\n",
    "    def clone(self):\n",
    "        \"\"\"Clone the full instance.\"\"\"\n",
    "        return pickle.loads(pickle.dumps(self))\n",
    "\n",
    "    def to(self, device):\n",
    "        \"\"\"Move to device.\"\"\"\n",
    "        if hasattr(self, \"c_repr\"):\n",
    "            self.c_repr = self.c_repr.to(device)\n",
    "        super().to(device)\n",
    "        return self\n",
    "\n",
    "    def __add__(self, model):\n",
    "        return SumEBM(self, model)\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"ConceptEBM\"}\n",
    "        model_dict[\"mode\"] = self.mode\n",
    "        model_dict[\"in_channels\"] = self.in_channels\n",
    "        model_dict[\"repr_dim\"] = self.repr_dim\n",
    "        model_dict[\"w_type\"] = self.w_type\n",
    "        model_dict[\"mask_mode\"] = self.mask_mode\n",
    "        model_dict[\"channel_base\"] = self.channel_base\n",
    "        model_dict[\"two_branch_mode\"] = self.two_branch_mode\n",
    "        model_dict[\"is_spec_norm\"] = self.is_spec_norm\n",
    "        model_dict[\"is_res\"] = self.is_res\n",
    "        model_dict[\"c_repr_mode\"] = self.c_repr_mode\n",
    "        model_dict[\"c_repr_first\"] = self.c_repr_first\n",
    "        model_dict[\"c_repr_base\"] = self.c_repr_base\n",
    "        model_dict[\"z_mode\"] = self.z_mode\n",
    "        model_dict[\"z_first\"] = self.z_first\n",
    "        model_dict[\"z_dim\"] = self.z_dim\n",
    "        model_dict[\"pos_embed_mode\"] = self.pos_embed_mode\n",
    "        model_dict[\"aggr_mode\"] = self.aggr_mode\n",
    "        model_dict[\"img_dims\"] = self.img_dims\n",
    "        model_dict[\"act_name\"] = self.act_name\n",
    "        model_dict[\"normalization_type\"] = self.normalization_type\n",
    "        model_dict[\"self_attn_mode\"] = self.self_attn_mode\n",
    "        model_dict[\"dropout\"] = self.dropout\n",
    "        model_dict[\"last_act_name\"] = self.last_act_name\n",
    "        model_dict[\"n_avg_pool\"] = self.n_avg_pool\n",
    "        if hasattr(self, \"c_repr\"):\n",
    "            model_dict[\"c_repr\"] = to_np_array(self.c_repr)\n",
    "        if hasattr(self, \"c_str\"):\n",
    "            model_dict[\"c_str\"] = self.c_str\n",
    "        model_dict[\"state_dict\"] = to_cpu(self.state_dict())\n",
    "        return model_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1.3 ConceptEBMLarge:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConceptEBMLarge(nn.Module):\n",
    "    \"\"\"\n",
    "    An EBM designed specifically to find Concepts, e.g. Rect, Line, RectSolid or\n",
    "        Relations, e.g. SameShape, SameColor.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(\n",
    "        self,\n",
    "        mode=\"concept\",\n",
    "        in_channels=10,\n",
    "        repr_dim=4,\n",
    "        w_type=\"image+mask\",\n",
    "        mask_mode=\"concat\",\n",
    "        channel_base=128,\n",
    "        two_branch_mode=\"concat\",\n",
    "        is_spec_norm=True,\n",
    "        is_res=False,\n",
    "        c_repr_mode=\"l1\",\n",
    "        c_repr_first=0,\n",
    "        c_repr_base=2,\n",
    "        z_mode=\"None\",\n",
    "        z_first=2,\n",
    "        z_dim=4,\n",
    "        pos_embed_mode=\"None\",\n",
    "        aggr_mode=\"sum\",\n",
    "        img_dims=2,\n",
    "        act_name=\"leakyrelu0.2\",\n",
    "        normalization_type=\"None\",\n",
    "        dropout=0,\n",
    "        self_attn_mode=\"None\",\n",
    "        last_act_name=\"None\",\n",
    "        is_multiscale=True,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Initialize the EBM.\n",
    "\n",
    "        Args:\n",
    "            w_type: type of the first two arities of input.\n",
    "                choose from \"image\", \"mask\", \"image+mask\", \"obj\", \"image+obj\"\n",
    "        \"\"\"\n",
    "        super(ConceptEBMLarge, self).__init__()\n",
    "        self.mode = mode\n",
    "        self.in_channels = in_channels\n",
    "        self.repr_dim = repr_dim\n",
    "        self.w_type = w_type\n",
    "        self.w_dim = 1 if \"mask\" in self.w_type else in_channels\n",
    "        self.mask_mode = mask_mode\n",
    "        self.channel_base = channel_base\n",
    "        self.two_branch_mode = two_branch_mode\n",
    "\n",
    "        if self.mode in [\"concept\"]:\n",
    "            self.mask_arity = 1\n",
    "        elif self.mode in [\"operator\"]:\n",
    "            self.mask_arity = 2\n",
    "        else:\n",
    "            raise\n",
    "        self.is_spec_norm = is_spec_norm\n",
    "        self.is_res = is_res\n",
    "        self.c_repr_mode = c_repr_mode\n",
    "        self.c_repr_first = c_repr_first\n",
    "        self.c_repr_base = c_repr_base\n",
    "        self.z_mode = z_mode\n",
    "        self.z_first = z_first\n",
    "        self.z_dim = z_dim\n",
    "        self.pos_embed_mode = pos_embed_mode\n",
    "        self.aggr_mode = aggr_mode\n",
    "        self.img_dims = img_dims\n",
    "        self.act_name = act_name\n",
    "        self.act = get_activation(act_name)\n",
    "        self.normalization_type = normalization_type\n",
    "        self.dropout = dropout\n",
    "        assert self_attn_mode in [\"None\", \"pixel\"]\n",
    "        self.self_attn_mode = self_attn_mode\n",
    "        self.last_act_name = last_act_name\n",
    "        self.is_multiscale = is_multiscale\n",
    "        if img_dims == 2:\n",
    "            self.kernel_size = 3\n",
    "            self.padding = 1\n",
    "        elif img_dims == 1:\n",
    "            self.kernel_size = (3, 1)\n",
    "            self.padding = (1, 0)\n",
    "        else:\n",
    "            raise\n",
    "        self.init_main_model()\n",
    "        if self.is_multiscale:\n",
    "            self.init_mid_model()\n",
    "            self.init_small_model()\n",
    "\n",
    "\n",
    "    def init_small_model(self):\n",
    "        mask_mode = self.mask_mode\n",
    "        channel_base = self.channel_base\n",
    "        repr_dim = self.repr_dim\n",
    "        is_spec_norm = self.is_spec_norm\n",
    "        is_res = self.is_res\n",
    "        c_repr_mode = self.c_repr_mode\n",
    "        c_repr_first = self.c_repr_first\n",
    "        c_repr_base = self.c_repr_base\n",
    "        z_mode = self.z_mode\n",
    "        z_dim = self.z_dim\n",
    "        z_first = self.z_first\n",
    "        img_dims = self.img_dims\n",
    "        act_name = self.act_name\n",
    "        normalization_type = self.normalization_type\n",
    "        dropout = self.dropout\n",
    "\n",
    "        if is_spec_norm in [True, \"True\"]:\n",
    "            self.small_conv1 = spectral_norm(nn.Conv2d(self.in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base if self.mode==\"concept\" else channel_base//2, self.kernel_size, padding=self.padding), std=1)\n",
    "        elif is_spec_norm in [False, \"False\"]:\n",
    "            self.small_conv1 = nn.Conv2d(self.in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base if self.mode==\"concept\" else channel_base//2, self.kernel_size, padding=self.padding)\n",
    "        elif is_spec_norm == \"ws\":\n",
    "            self.small_conv1 = WSConv2d(self.in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base if self.mode==\"concept\" else channel_base//2, self.kernel_size, padding=self.padding)\n",
    "        else:\n",
    "            raise\n",
    "        \n",
    "        self.small_res_1a = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "        self.small_res_1b = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "\n",
    "        self.small_res_2a = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "        self.small_res_2b = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv+rescale\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "\n",
    "        self.small_energy_map = nn.Linear(channel_base*2, 1)\n",
    "\n",
    "\n",
    "    def small_model(self, input, mask, c_repr=None, z=None):\n",
    "        c_repr_first_dict = {0: 0, 1: 2, 2: 4, 3: 6}\n",
    "        z_first_dict = {0: 0, 1: 2, 2: 4, 3: 6}\n",
    "        if self.mode == \"concept\":\n",
    "            assert len(mask) == 1\n",
    "            if self.mask_mode == \"concat\":\n",
    "                input_aug = torch.cat([input, mask[0]], 1)  # input: [B, C, H, W], mask: [B, c, H, W]\n",
    "            elif self.mask_mode == \"mul\":\n",
    "                input_aug = input * mask[0]\n",
    "            elif self.mask_mode == \"mulcat\":\n",
    "                input_aug = torch.cat([input*mask[0], mask[0]], 1)\n",
    "            elif self.mask_mode == \"fil\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_aug = torch.cat([torch.clamp(input[:,:1]+(1-mask[0]), 0, 1), input[:,1:]*mask[0]], 1)\n",
    "            elif self.mask_mode == \"filcat\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_fil = torch.cat([torch.clamp(input[:,:1]+(1-mask[0]), 0, 1), input[:,1:]*mask[0]], 1)\n",
    "                input_aug = torch.cat([input_fil, mask[0]], 1)\n",
    "            else:\n",
    "                raise\n",
    "            x = F.avg_pool2d(input_aug, 3, stride=2, padding=1)\n",
    "            x = F.avg_pool2d(x, 3, stride=2, padding=1)\n",
    "            x = self.act(self.small_conv1(x))\n",
    "\n",
    "        elif self.mode == \"operator\":\n",
    "            # Combining input image and mask:\n",
    "            if not (isinstance(input, tuple) or isinstance(input, list)):\n",
    "                input = (input, input)\n",
    "            if self.mask_mode == \"concat\":\n",
    "                input_aug_0 = torch.cat([input[0], mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([input[1], mask[1]], 1)\n",
    "            elif self.mask_mode == \"mul\":\n",
    "                input_aug_0 = input[0] * mask[0]\n",
    "                input_aug_1 = input[1] * mask[1]\n",
    "            elif self.mask_mode == \"mulcat\":\n",
    "                input_aug_0 = torch.cat([input[0]*mask[0], mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([input[1]*mask[1], mask[1]], 1)\n",
    "            elif self.mask_mode == \"fil\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_aug_0 = torch.cat([torch.clamp(input[0][:,:1]+(1-mask[0]), 0, 1), input[0][:,1:]*mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([torch.clamp(input[1][:,:1]+(1-mask[1]), 0, 1), input[1][:,1:]*mask[1]], 1)\n",
    "            elif self.mask_mode == \"filcat\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_fil_0 = torch.cat([torch.clamp(input[0][:,:1]+(1-mask[0]), 0, 1), input[0][:,1:]*mask[0]], 1)\n",
    "                input_aug_0 = torch.cat([input_fil_0, mask[0]], 1)\n",
    "                input_fil_1 = torch.cat([torch.clamp(input[1][:,:1]+(1-mask[1]), 0, 1), input[1][:,1:]*mask[1]], 1)\n",
    "                input_aug_1 = torch.cat([input_fil_1, mask[1]], 1)\n",
    "            else:\n",
    "                raise\n",
    "\n",
    "            x_0 = F.avg_pool2d(input_aug_0, 3, stride=2, padding=1)\n",
    "            x_0 = F.avg_pool2d(x_0, 3, stride=2, padding=1)\n",
    "            x_0 = self.act(self.small_conv1(x_0))\n",
    "            x_1 = F.avg_pool2d(input_aug_1, 3, stride=2, padding=1)\n",
    "            x_1 = F.avg_pool2d(x_1, 3, stride=2, padding=1)\n",
    "            x_1 = self.act(self.small_conv1(x_1))\n",
    "            x = torch.cat([x_0, x_1], 1)\n",
    "\n",
    "        x = self.small_res_1a(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=0 else None, z=z if z_first_dict[self.z_first]<=0 else None)\n",
    "        x = self.small_res_1b(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=1 else None, z=z if z_first_dict[self.z_first]<=1 else None)\n",
    "\n",
    "        x = self.small_res_2a(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=2 else None, z=z if z_first_dict[self.z_first]<=2 else None)\n",
    "        x = self.small_res_2b(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=3 else None, z=z if z_first_dict[self.z_first]<=3 else None)\n",
    "        x = self.act(x)\n",
    "\n",
    "        if self.aggr_mode == \"sum\":\n",
    "            x = x.view(x.shape[0], x.shape[1], -1).sum(2)\n",
    "        elif self.aggr_mode == \"max\":\n",
    "            x = x.view(x.shape[0], x.shape[1], -1).max(2)[0]\n",
    "        elif self.aggr_mode == \"mean\":\n",
    "            x = x.view(x.shape[0], x.shape[1], -1).mean(2)\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        energy = self.small_energy_map(x)\n",
    "\n",
    "        # Last activation:\n",
    "        if self.last_act_name != \"None\":\n",
    "            if self.last_act_name in [\"square\", \"softplus\", \"exp\", \"sigmoid\"]:\n",
    "                energy = get_activation(self.last_act_name)(energy)\n",
    "            else:\n",
    "                raise\n",
    "        return energy\n",
    "\n",
    "\n",
    "    def init_mid_model(self):\n",
    "        mask_mode = self.mask_mode\n",
    "        channel_base = self.channel_base\n",
    "        repr_dim = self.repr_dim\n",
    "        is_spec_norm = self.is_spec_norm\n",
    "        is_res = self.is_res\n",
    "        c_repr_mode = self.c_repr_mode\n",
    "        c_repr_first = self.c_repr_first\n",
    "        c_repr_base = self.c_repr_base\n",
    "        z_mode = self.z_mode\n",
    "        z_dim = self.z_dim\n",
    "        z_first = self.z_first\n",
    "        img_dims = self.img_dims\n",
    "        act_name = self.act_name\n",
    "        normalization_type = self.normalization_type\n",
    "        dropout = self.dropout\n",
    "\n",
    "        if is_spec_norm in [True, \"True\"]:\n",
    "            self.mid_conv1 = spectral_norm(nn.Conv2d(self.in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base if self.mode==\"concept\" else channel_base//2, self.kernel_size, padding=self.padding), std=1)\n",
    "        elif is_spec_norm in [False, \"False\"]:\n",
    "            self.mid_conv1 = nn.Conv2d(self.in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base if self.mode==\"concept\" else channel_base//2, self.kernel_size, padding=self.padding)\n",
    "        elif is_spec_norm == \"ws\":\n",
    "            self.mid_conv1 = WSConv2d(self.in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base if self.mode==\"concept\" else channel_base//2, self.kernel_size, padding=self.padding)\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        self.mid_res_1a = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "        self.mid_res_1b = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "\n",
    "        self.mid_res_2a = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "        self.mid_res_2b = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv+rescale\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "\n",
    "        self.mid_res_3a = CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, downsample=False, is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=2 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=2 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "        self.mid_res_3b = CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv+rescale\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=2 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=2 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "\n",
    "        self.mid_energy_map = nn.Linear(channel_base*4, 1)\n",
    "        self.avg_pool = Downsample(channels=3)\n",
    "\n",
    "\n",
    "    def mid_model(self, input, mask, c_repr=None, z=None):\n",
    "        c_repr_first_dict = {0: 0, 1: 2, 2: 4, 3: 6}\n",
    "        z_first_dict = {0: 0, 1: 2, 2: 4, 3: 6}\n",
    "        if self.mode == \"concept\":\n",
    "            assert len(mask) == 1\n",
    "            if self.mask_mode == \"concat\":\n",
    "                input_aug = torch.cat([input, mask[0]], 1)  # input: [B, C, H, W], mask: [B, c, H, W]\n",
    "            elif self.mask_mode == \"mul\":\n",
    "                input_aug = input * mask[0]\n",
    "            elif self.mask_mode == \"mulcat\":\n",
    "                input_aug = torch.cat([input*mask[0], mask[0]], 1)\n",
    "            elif self.mask_mode == \"fil\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_aug = torch.cat([torch.clamp(input[:,:1]+(1-mask[0]), 0, 1), input[:,1:]*mask[0]], 1)\n",
    "            elif self.mask_mode == \"filcat\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_fil = torch.cat([torch.clamp(input[:,:1]+(1-mask[0]), 0, 1), input[:,1:]*mask[0]], 1)\n",
    "                input_aug = torch.cat([input_fil, mask[0]], 1)\n",
    "            else:\n",
    "                raise\n",
    "            x = F.avg_pool2d(input_aug, 3, stride=2, padding=1)\n",
    "            x = self.act(self.mid_conv1(x))\n",
    "        \n",
    "        elif self.mode == \"operator\":\n",
    "            # Combining input image and mask:\n",
    "            if not (isinstance(input, tuple) or isinstance(input, list)):\n",
    "                input = (input, input)\n",
    "            if self.mask_mode == \"concat\":\n",
    "                input_aug_0 = torch.cat([input[0], mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([input[1], mask[1]], 1)\n",
    "            elif self.mask_mode == \"mul\":\n",
    "                input_aug_0 = input[0] * mask[0]\n",
    "                input_aug_1 = input[1] * mask[1]\n",
    "            elif self.mask_mode == \"mulcat\":\n",
    "                input_aug_0 = torch.cat([input[0]*mask[0], mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([input[1]*mask[1], mask[1]], 1)\n",
    "            elif self.mask_mode == \"fil\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_aug_0 = torch.cat([torch.clamp(input[0][:,:1]+(1-mask[0]), 0, 1), input[0][:,1:]*mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([torch.clamp(input[1][:,:1]+(1-mask[1]), 0, 1), input[1][:,1:]*mask[1]], 1)\n",
    "            elif self.mask_mode == \"filcat\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_fil_0 = torch.cat([torch.clamp(input[0][:,:1]+(1-mask[0]), 0, 1), input[0][:,1:]*mask[0]], 1)\n",
    "                input_aug_0 = torch.cat([input_fil_0, mask[0]], 1)\n",
    "                input_fil_1 = torch.cat([torch.clamp(input[1][:,:1]+(1-mask[1]), 0, 1), input[1][:,1:]*mask[1]], 1)\n",
    "                input_aug_1 = torch.cat([input_fil_1, mask[1]], 1)\n",
    "            else:\n",
    "                raise\n",
    "\n",
    "            x_0 = F.avg_pool2d(input_aug_0, 3, stride=2, padding=1)\n",
    "            x_0 = self.act(self.mid_conv1(x_0))\n",
    "            x_1 = F.avg_pool2d(input_aug_1, 3, stride=2, padding=1)\n",
    "            x_1 = self.act(self.mid_conv1(x_1))\n",
    "            x = torch.cat([x_0, x_1], 1)\n",
    "\n",
    "        x = self.mid_res_1a(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=0 else None, z=z if z_first_dict[self.z_first]<=0 else None)\n",
    "        x = self.mid_res_1b(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=1 else None, z=z if z_first_dict[self.z_first]<=1 else None)\n",
    "\n",
    "        x = self.mid_res_2a(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=2 else None, z=z if z_first_dict[self.z_first]<=2 else None)\n",
    "        x = self.mid_res_2b(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=3 else None, z=z if z_first_dict[self.z_first]<=3 else None)\n",
    "\n",
    "        x = self.mid_res_3a(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=4 else None, z=z if z_first_dict[self.z_first]<=4 else None)\n",
    "        x = self.mid_res_3b(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=5 else None, z=z if z_first_dict[self.z_first]<=5 else None)\n",
    "        x = self.act(x)\n",
    "\n",
    "        if self.aggr_mode == \"sum\":\n",
    "            x = x.view(x.shape[0], x.shape[1], -1).sum(2)\n",
    "        elif self.aggr_mode == \"max\":\n",
    "            x = x.view(x.shape[0], x.shape[1], -1).max(2)[0]\n",
    "        elif self.aggr_mode == \"mean\":\n",
    "            x = x.view(x.shape[0], x.shape[1], -1).mean(2)\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        energy = self.mid_energy_map(x)\n",
    "\n",
    "        # Last activation:\n",
    "        if self.last_act_name != \"None\":\n",
    "            if self.last_act_name in [\"square\", \"softplus\", \"exp\", \"sigmoid\"]:\n",
    "                energy = get_activation(self.last_act_name)(energy)\n",
    "            else:\n",
    "                raise\n",
    "        return energy\n",
    "\n",
    "\n",
    "    def init_main_model(self):\n",
    "        mask_mode = self.mask_mode\n",
    "        channel_base = self.channel_base\n",
    "        repr_dim = self.repr_dim\n",
    "        is_spec_norm = self.is_spec_norm\n",
    "        is_res = self.is_res\n",
    "        c_repr_mode = self.c_repr_mode\n",
    "        c_repr_first = self.c_repr_first\n",
    "        c_repr_base = self.c_repr_base\n",
    "        z_mode = self.z_mode\n",
    "        z_dim = self.z_dim\n",
    "        z_first = self.z_first\n",
    "        img_dims = self.img_dims\n",
    "        act_name = self.act_name\n",
    "        normalization_type = self.normalization_type\n",
    "        dropout = self.dropout\n",
    "\n",
    "        if is_spec_norm in [True, \"True\"]:\n",
    "            self.conv1 = spectral_norm(nn.Conv2d(self.in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base//2 if self.mode==\"concept\" else channel_base//4, self.kernel_size, padding=self.padding), std=1)\n",
    "        elif is_spec_norm in [False, \"False\"]:\n",
    "            self.conv1 = nn.Conv2d(self.in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base//2 if self.mode==\"concept\" else channel_base//4, self.kernel_size, padding=self.padding)\n",
    "        elif is_spec_norm == \"ws\":\n",
    "            self.conv1 = WSConv2d(self.in_channels+self.w_dim if mask_mode in [\"concat\", \"mulcat\", \"filcat\"] else in_channels, channel_base//2 if self.mode==\"concept\" else channel_base//4, self.kernel_size, padding=self.padding)\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        self.res_1a = CResBlock(channel_base//2, channel_base//2, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv+rescale\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "        self.res_1b = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=0 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=0 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "\n",
    "        self.res_2a = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "        self.res_2b = CResBlock(channel_base, channel_base, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv+rescale\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=1 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=1 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "\n",
    "        self.res_3a = CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, downsample=False, is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=2 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=2 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "        self.res_3b = CResBlock(channel_base*2, channel_base*2, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv+rescale\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=2 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=2 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "\n",
    "        self.res_4a = CResBlock(channel_base*4, channel_base*4, repr_dim=repr_dim, downsample=False, is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=3 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=3 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "        self.res_4b = CResBlock(channel_base*4, channel_base*4, repr_dim=repr_dim, downsample=True, downsample_mode=\"conv+rescale\", is_res=is_res, is_spec_norm=is_spec_norm, c_repr_mode=c_repr_mode if c_repr_first<=3 else \"None\", c_repr_base=c_repr_base, z_mode=z_mode if z_first<=3 else \"None\", z_dim=z_dim, img_dims=img_dims, act_name=act_name, normalization_type=normalization_type, dropout=dropout)\n",
    "\n",
    "        if self.self_attn_mode in [\"pixel\"]:\n",
    "            self.self_attn = Self_Attn(4 * channel_base, self.act_name)\n",
    "\n",
    "        self.energy_map = nn.Linear(channel_base*8, 1)\n",
    "\n",
    "\n",
    "    def main_model(self, input, mask, c_repr=None, z=None):\n",
    "        c_repr_first_dict = {0: 0, 1: 2, 2: 4, 3: 6}\n",
    "        z_first_dict = {0: 0, 1: 2, 2: 4, 3: 6}\n",
    "        if self.mode == \"concept\":\n",
    "            assert len(mask) == 1\n",
    "            if self.mask_mode == \"concat\":\n",
    "                input_aug = torch.cat([input, mask[0]], 1)  # input: [B, C, H, W], mask: [B, c, H, W]\n",
    "            elif self.mask_mode == \"mul\":\n",
    "                input_aug = input * mask[0]\n",
    "            elif self.mask_mode == \"mulcat\":\n",
    "                input_aug = torch.cat([input*mask[0], mask[0]], 1)\n",
    "            elif self.mask_mode == \"fil\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_aug = torch.cat([torch.clamp(input[:,:1]+(1-mask[0]), 0, 1), input[:,1:]*mask[0]], 1)\n",
    "            elif self.mask_mode == \"filcat\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_fil = torch.cat([torch.clamp(input[:,:1]+(1-mask[0]), 0, 1), input[:,1:]*mask[0]], 1)\n",
    "                input_aug = torch.cat([input_fil, mask[0]], 1)\n",
    "            else:\n",
    "                raise\n",
    "            x = self.act(self.conv1(input_aug))\n",
    "\n",
    "        elif self.mode == \"operator\":\n",
    "            # Combining input image and mask:\n",
    "            if not (isinstance(input, tuple) or isinstance(input, list)):\n",
    "                input = (input, input)\n",
    "            if self.mask_mode == \"concat\":\n",
    "                input_aug_0 = torch.cat([input[0], mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([input[1], mask[1]], 1)\n",
    "            elif self.mask_mode == \"mul\":\n",
    "                input_aug_0 = input[0] * mask[0]\n",
    "                input_aug_1 = input[1] * mask[1]\n",
    "            elif self.mask_mode == \"mulcat\":\n",
    "                input_aug_0 = torch.cat([input[0]*mask[0], mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([input[1]*mask[1], mask[1]], 1)\n",
    "            elif self.mask_mode == \"fil\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_aug_0 = torch.cat([torch.clamp(input[0][:,:1]+(1-mask[0]), 0, 1), input[0][:,1:]*mask[0]], 1)\n",
    "                input_aug_1 = torch.cat([torch.clamp(input[1][:,:1]+(1-mask[1]), 0, 1), input[1][:,1:]*mask[1]], 1)\n",
    "            elif self.mask_mode == \"filcat\":\n",
    "                assert \"obj\" not in self.w_type\n",
    "                input_fil_0 = torch.cat([torch.clamp(input[0][:,:1]+(1-mask[0]), 0, 1), input[0][:,1:]*mask[0]], 1)\n",
    "                input_aug_0 = torch.cat([input_fil_0, mask[0]], 1)\n",
    "                input_fil_1 = torch.cat([torch.clamp(input[1][:,:1]+(1-mask[1]), 0, 1), input[1][:,1:]*mask[1]], 1)\n",
    "                input_aug_1 = torch.cat([input_fil_1, mask[1]], 1)\n",
    "            else:\n",
    "                raise\n",
    "\n",
    "            x_0 = self.act(self.conv1(input_aug_0))\n",
    "            x_1 = self.act(self.conv1(input_aug_1))\n",
    "            x = torch.cat([x_0, x_1], 1)\n",
    "\n",
    "        x = self.res_1a(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=0 else None, z=z if z_first_dict[self.z_first]<=0 else None)\n",
    "        x = self.res_1b(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=1 else None, z=z if z_first_dict[self.z_first]<=1 else None)\n",
    "\n",
    "        x = self.res_2a(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=2 else None, z=z if z_first_dict[self.z_first]<=2 else None)\n",
    "        x = self.res_2b(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=3 else None, z=z if z_first_dict[self.z_first]<=3 else None)\n",
    "\n",
    "        x = self.res_3a(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=4 else None, z=z if z_first_dict[self.z_first]<=4 else None)\n",
    "        x = self.res_3b(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=5 else None, z=z if z_first_dict[self.z_first]<=5 else None)\n",
    "\n",
    "        if self.self_attn_mode in [\"pixel\"]:\n",
    "            x, _ = self.self_attn(x)\n",
    "\n",
    "        x = self.res_4a(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=6 else None, z=z if z_first_dict[self.z_first]<=6 else None)\n",
    "        x = self.res_4b(x, c_repr=c_repr if c_repr_first_dict[self.c_repr_first]<=7 else None, z=z if z_first_dict[self.z_first]<=7 else None)\n",
    "        x = self.act(x)\n",
    "\n",
    "        if self.aggr_mode == \"sum\":\n",
    "            x = x.view(x.shape[0], x.shape[1], -1).sum(2)\n",
    "        elif self.aggr_mode == \"max\":\n",
    "            x = x.view(x.shape[0], x.shape[1], -1).max(2)[0]\n",
    "        elif self.aggr_mode == \"mean\":\n",
    "            x = x.view(x.shape[0], x.shape[1], -1).mean(2)\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        energy = self.energy_map(x)\n",
    "\n",
    "        # Last activation:\n",
    "        if self.last_act_name != \"None\":\n",
    "            if self.last_act_name in [\"square\", \"softplus\", \"exp\", \"sigmoid\"]:\n",
    "                energy = get_activation(self.last_act_name)(energy)\n",
    "            else:\n",
    "                raise\n",
    "        return energy\n",
    "\n",
    "\n",
    "    def forward(self, input, mask, c_repr=None, z=None, **kwargs):\n",
    "        energy = self.main_model(input, mask, c_repr=c_repr, z=z)\n",
    "\n",
    "        if self.is_multiscale:\n",
    "            large_energy = energy\n",
    "            mid_energy = self.mid_model(input, mask, c_repr=c_repr, z=z)\n",
    "            small_energy = self.small_model(input, mask, c_repr=c_repr, z=z)\n",
    "            energy = torch.cat([small_energy, mid_energy, large_energy], dim=-1).sum(-1, keepdims=True)\n",
    "        return energy\n",
    "\n",
    "\n",
    "    def classify(self, input, mask, concept_collection, topk=-1, CONCEPTS=None, OPERATORS=None):\n",
    "        \"\"\"\n",
    "        Given the input and mask, classify the selected concept by picking the\n",
    "        lowest-energy concept from concept_collection.\n",
    "        \"\"\"\n",
    "        if isinstance(input, tuple) or isinstance(input, list):\n",
    "            length = len(input[0])\n",
    "            device = input[0].device\n",
    "        else:\n",
    "            length = len(input)\n",
    "            device = input.device\n",
    "        if topk == -1:\n",
    "            topk = len(concept_collection)\n",
    "        c_repr_energy = []\n",
    "        for j in range(len(concept_collection)):\n",
    "            c_repr = id_to_tensor([concept_collection[j]] * length, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device) # len 4\n",
    "            neg_energy = self(input, mask=mask, c_repr=c_repr)\n",
    "            c_repr_energy.append(neg_energy)\n",
    "        c_repr_energy = torch.cat(c_repr_energy, 1)\n",
    "        c_repr_argsort = c_repr_energy.argsort(1)\n",
    "        c_repr_pred_list = []\n",
    "        for i, argsort in enumerate(c_repr_argsort):\n",
    "            c_repr_pred = {}\n",
    "            for k in range(min(topk, len(concept_collection))):\n",
    "                id_k = c_repr_argsort[i][k]\n",
    "                c_repr_pred[concept_collection[id_k]] = c_repr_energy[i][id_k].item()\n",
    "            c_repr_pred_list.append(c_repr_pred)\n",
    "        return c_repr_pred_list\n",
    "\n",
    "\n",
    "    def ground(\n",
    "        self,\n",
    "        input,\n",
    "        args,\n",
    "        mask=None,\n",
    "        c_repr=None,\n",
    "        z=None,\n",
    "        ensemble_size=18,\n",
    "        topk=-1,\n",
    "        w_init_type=\"random\",\n",
    "        sample_step=150,\n",
    "        is_grad=False,\n",
    "        is_return_E=False,\n",
    "        isplot=True,\n",
    "        **kwargs\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Given an input image, find the best mask to match the given (optional)\n",
    "        concept representation.\n",
    "\n",
    "        Internally, uses a model ensemble for best results.\n",
    "        \"\"\"\n",
    "        \n",
    "        def init_neg_mask(input, init, ensemble_size):\n",
    "            \"\"\"Initialize negative mask\"\"\"\n",
    "            if isinstance(input, tuple):\n",
    "                assert len(input[0].shape) == len(input[1].shape) == 4\n",
    "                device = input[0].device\n",
    "                w_dim = 1 if \"mask\" in self.w_type else input[0].shape[1]\n",
    "                neg_mask = (torch.rand(input[0].shape[0]*ensemble_size, w_dim, *input[0].shape[2:]).to(device), torch.rand(input[1].shape[0]*ensemble_size, w_dim, *input[1].shape[2:]).to(device))\n",
    "                if init == \"input-mask\":\n",
    "                    assert input[0].shape[1] == 10\n",
    "                    input_l = repeat_n(input, n_repeats=ensemble_size)\n",
    "                    neg_mask = (neg_mask[0] * (input_l[0].argmax(1)[:, None] != 0), neg_mask[1] * (input_l[1].argmax(1)[:, None] != 0))\n",
    "                neg_mask[0].requires_grad = True\n",
    "                neg_mask[1].requires_grad = True\n",
    "            else:\n",
    "                assert len(input.shape) == 4\n",
    "                device = input.device\n",
    "                w_dim = 1 if \"mask\" in self.w_type else input.shape[1]\n",
    "                neg_mask = (torch.rand(input.shape[0]*ensemble_size, w_dim, *input.shape[2:]).to(device),)\n",
    "                if init == \"input-mask\":\n",
    "                    assert input.shape[1] == 10\n",
    "                    input_l = repeat_n(input, n_repeats=ensemble_size)\n",
    "                    neg_mask = (neg_mask[0] * (input_l.argmax(1)[:, None] != 0),)\n",
    "                neg_mask[0].requires_grad = True\n",
    "            return neg_mask\n",
    "\n",
    "        # Update args:\n",
    "        args = deepcopy(args)\n",
    "        for key, value in kwargs.items():\n",
    "            setattr(args, key, value)\n",
    "        args.sample_step = sample_step\n",
    "        if isinstance(input, tuple) or isinstance(input, list):\n",
    "            args.is_image_tuple = True\n",
    "            device = input[0].device\n",
    "        else:\n",
    "            args.is_image_tuple = False\n",
    "            device = input.device\n",
    "\n",
    "        # Perform SGLD:\n",
    "        if mask is None:\n",
    "            neg_mask = init_neg_mask(input, init=w_init_type, ensemble_size=ensemble_size)\n",
    "        else:\n",
    "            neg_mask = tuple([repeat_n(mask[0], n_repeats=ensemble_size)])\n",
    "\n",
    "        if self.z_mode != \"None\":\n",
    "            if z is None:\n",
    "                z = tuple([torch.rand(neg_mask[0].shape[0], self.z_dim, device=device)])\n",
    "            else:\n",
    "                z = tuple([repeat_n(z[0], n_repeats=ensemble_size)])\n",
    "\n",
    "        (img_ensemble, neg_mask_ensemble, z_ensemble, zgnn_ensemble, wtarget_ensemble), neg_out_list_ensemble, info_ensemble = neg_mask_sgd_ensemble(\n",
    "            self, input, neg_mask, c_repr, z=z, zgnn=None, wtarget=None, args=args,\n",
    "            ensemble_size=ensemble_size, is_grad=is_grad,\n",
    "            is_return_E=is_return_E,\n",
    "        )\n",
    "\n",
    "        neg_out_ensemble = neg_out_list_ensemble[-1]  # neg_out_ensemble: [ensemble_size, B]\n",
    "        # Sort the obtained results by energy for each example:\n",
    "        neg_out_ensemble = torch.FloatTensor(neg_out_ensemble).transpose(0,1)  # [B, ensemble_size]\n",
    "        neg_out_argsort = neg_out_ensemble.argsort(1)  # [B, ensemble_size]\n",
    "        batch_size = neg_out_argsort.shape[0]\n",
    "        neg_out_ensemble_sorted = torch.stack([neg_out_ensemble[i][neg_out_argsort[i]] for i in range(batch_size)])\n",
    "        if zgnn_ensemble is not None or wtarget_ensemble is not None:\n",
    "            neg_task_out_ensemble = neg_out_ensemble.reshape(*batch_shape, -1).mean(1)  # [B_task, ensemble_size]\n",
    "            neg_task_out_argsort = neg_task_out_ensemble.argsort(1)\n",
    "\n",
    "        if img_ensemble is not None:\n",
    "            if args.is_image_tuple:\n",
    "                img_ensemble = tuple(img_ensemble[k].transpose(0,1) for k in range(len(img)))  # Each element [B, ensemble_size, C, H, W]\n",
    "                img_ensemble_sorted = []\n",
    "                for k in range(len(img)):\n",
    "                    img_ensemble_sorted.append(torch.stack([img_ensemble[k][i][neg_out_argsort[i]] for i in range(batch_size)]))\n",
    "                img_ensemble_sorted = tuple(img_ensemble_sorted)  # each element: [B, ensemble_size, C, H, W] sorted along dim=1 according to neg_out\n",
    "            else:\n",
    "                img_ensemble = img_ensemble.transpose(0,1)  # [B, ensemble_size, C, H, W]\n",
    "                img_ensemble_sorted = torch.stack([img_ensemble[i][neg_out_argsort[i]] for i in range(batch_size)])\n",
    "        else:\n",
    "            img_ensemble_sorted = None\n",
    "\n",
    "        if neg_mask_ensemble is not None:\n",
    "            neg_mask_ensemble = tuple(neg_mask_ensemble[k].transpose(0,1) for k in range(self.mask_arity))  # Each element [B, ensemble_size, C, H, W]\n",
    "            neg_mask_ensemble_sorted = []\n",
    "            for k in range(self.mask_arity):\n",
    "                neg_mask_ensemble_sorted.append(torch.stack([neg_mask_ensemble[k][i][neg_out_argsort[i]] for i in range(len(neg_mask_ensemble[0]))]))\n",
    "            neg_mask_ensemble_sorted = tuple(neg_mask_ensemble_sorted)  # each element: [B, ensemble_size, C, H, W] sorted along dim=1 according to neg_out\n",
    "        else:\n",
    "            neg_mask_ensemble_sorted = None\n",
    "\n",
    "        if z_ensemble is not None:\n",
    "            z_ensemble = tuple(z_ensemble[k].transpose(0,1) for k in range(len(z_ensemble)))  # Each element [B, ensemble_size, Z]\n",
    "            z_ensemble_sorted = []\n",
    "            for k in range(len(z_ensemble)):\n",
    "                z_ensemble_sorted.append(torch.stack([z_ensemble[k][i][neg_out_argsort[i]] for i in range(batch_size)]))\n",
    "            z_ensemble_sorted = tuple(z_ensemble_sorted)  # each element: [B, ensemble_size, Z] sorted along dim=1 according to neg_out\n",
    "        else:\n",
    "            z_ensemble_sorted = None\n",
    "\n",
    "        if zgnn_ensemble is not None:\n",
    "            zgnn_ensemble = tuple(zgnn_ensemble[k].transpose(0,1) if zgnn_ensemble[k] is not None else None for k in range(len(zgnn_ensemble)))  # Each element [B, ensemble_size, Zgnn]\n",
    "            zgnn_ensemble_sorted = []\n",
    "            for k in range(len(zgnn_ensemble)):\n",
    "                if zgnn_ensemble[k] is not None:\n",
    "                    zgnn_ensemble_sorted.append(torch.stack([zgnn_ensemble[k][i][neg_task_out_argsort[i]][:topk_core] for i in range(batch_shape[0])]))\n",
    "                else:\n",
    "                    zgnn_ensemble_sorted.append(None)\n",
    "            zgnn_ensemble_sorted = tuple(zgnn_ensemble_sorted)  # each element: [B, ensemble_size, zgnn_dim] sorted along dim=1 according to neg_out\n",
    "        else:\n",
    "            zgnn_ensemble_sorted = None\n",
    "\n",
    "        if wtarget_ensemble is not None:\n",
    "            wtarget_ensemble = wtarget_ensemble.transpose(0,1)  # [B, ensemble_size, w_dim, H, W]\n",
    "            wtarget_ensemble_sorted = torch.stack([wtarget_ensemble[i][neg_out_argsort[i]][:topk_core] for i in range(batch_size)])\n",
    "        else:\n",
    "            wtarget_ensemble_sorted = None\n",
    "\n",
    "        if isplot >= 2:\n",
    "            # Plot SGLD learning curve:\n",
    "            print(\"SGLD learning curve:\")\n",
    "            plt.figure(figsize=(12,6))\n",
    "            for i in range(min(neg_out_list_ensemble.shape[-1], 6)):  # neg_out_list_ensemble: [sample_step, ensemble_size, B]\n",
    "#                 print(\"Example {}\".format(i))\n",
    "                for k in range(min(6, ensemble_size)):\n",
    "                    plt.plot(neg_out_list_ensemble[:, neg_out_argsort[i][k],i], c=COLOR_LIST[k], label=\"id_{}\".format(k), alpha=0.4)\n",
    "            plt.legend()\n",
    "            plt.show()\n",
    "        return (img_ensemble_sorted, neg_mask_ensemble_sorted, z_ensemble_sorted, zgnn_ensemble_sorted, wtarget_ensemble_sorted), neg_out_ensemble_sorted\n",
    "\n",
    "    def set_c(self, c_repr, c_str=None):\n",
    "        \"\"\"Set default c_repr.\"\"\"\n",
    "        assert len(c_repr.shape) == 2 and c_repr.shape[0] == 1\n",
    "        if not isinstance(c_repr, torch.Tensor):\n",
    "            c_repr = torch.FloatTensor(c_repr)\n",
    "        device = next(iter(self.parameters())).device\n",
    "        self.c_repr = c_repr.to(device)\n",
    "        self.c_str = c_str\n",
    "        return self\n",
    "\n",
    "    def set_repr_dict(self, c_repr_dict):\n",
    "        \"\"\"Set the dictionary of c_repr.\"\"\"\n",
    "        self.c_repr_dict = {}\n",
    "        device = next(iter(self.parameters())).device\n",
    "        for key, c_repr in c_repr_dict.items():\n",
    "            if not isinstance(c_repr, torch.Tensor):\n",
    "                c_repr = torch.FloatTensor(c_repr)\n",
    "            self.c_repr_dict[key] = c_repr.to(device)\n",
    "        return self\n",
    "\n",
    "    def clone(self):\n",
    "        \"\"\"Clone the full instance.\"\"\"\n",
    "        return pickle.loads(pickle.dumps(self))\n",
    "\n",
    "    def to(self, device):\n",
    "        \"\"\"Move to device.\"\"\"\n",
    "        if hasattr(self, \"c_repr\"):\n",
    "            self.c_repr = self.c_repr.to(device)\n",
    "        super().to(device)\n",
    "        return self\n",
    "\n",
    "    def __add__(self, model):\n",
    "        return SumEBM(self, model)\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"ConceptEBMLarge\"}\n",
    "        model_dict[\"mode\"] = self.mode\n",
    "        model_dict[\"in_channels\"] = self.in_channels\n",
    "        model_dict[\"repr_dim\"] = self.repr_dim\n",
    "        model_dict[\"w_type\"] = self.w_type\n",
    "        model_dict[\"mask_mode\"] = self.mask_mode\n",
    "        model_dict[\"channel_base\"] = self.channel_base\n",
    "        model_dict[\"two_branch_mode\"] = self.two_branch_mode\n",
    "        model_dict[\"is_spec_norm\"] = self.is_spec_norm\n",
    "        model_dict[\"is_res\"] = self.is_res\n",
    "        model_dict[\"c_repr_mode\"] = self.c_repr_mode\n",
    "        model_dict[\"c_repr_first\"] = self.c_repr_first\n",
    "        model_dict[\"c_repr_base\"] = self.c_repr_base\n",
    "        model_dict[\"z_mode\"] = self.z_mode\n",
    "        model_dict[\"z_first\"] = self.z_first\n",
    "        model_dict[\"z_dim\"] = self.z_dim\n",
    "        model_dict[\"pos_embed_mode\"] = self.pos_embed_mode\n",
    "        model_dict[\"aggr_mode\"] = self.aggr_mode\n",
    "        model_dict[\"img_dims\"] = self.img_dims\n",
    "        model_dict[\"act_name\"] = self.act_name\n",
    "        model_dict[\"normalization_type\"] = self.normalization_type\n",
    "        model_dict[\"self_attn_mode\"] = self.self_attn_mode\n",
    "        model_dict[\"dropout\"] = self.dropout\n",
    "        model_dict[\"last_act_name\"] = self.last_act_name\n",
    "        model_dict[\"is_multiscale\"] = self.is_multiscale\n",
    "        if hasattr(self, \"c_repr\"):\n",
    "            model_dict[\"c_repr\"] = to_np_array(self.c_repr)\n",
    "        if hasattr(self, \"c_str\"):\n",
    "            model_dict[\"c_str\"] = self.c_str\n",
    "        model_dict[\"state_dict\"] = to_cpu(self.state_dict())\n",
    "        return model_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CelebAModel(nn.Module):\n",
    "    \"\"\"From https://github.com/yilundu/improved_contrastive_divergence/blob/master/models.py#L413\"\"\"\n",
    "    def __init__(self, args, debug=False):\n",
    "        from improved_contrastive_divergence.models import CondResBlock\n",
    "        super(CelebAModel, self).__init__()\n",
    "        self.act = swish\n",
    "        self.max_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "        self.cond = args.cond\n",
    "\n",
    "        self.args = args\n",
    "        self.init_main_model()\n",
    "\n",
    "        if args.multiscale:\n",
    "            self.init_mid_model()\n",
    "            self.init_small_model()\n",
    "\n",
    "        self.relu = torch.nn.ReLU(inplace=True)\n",
    "        self.downsample = Downsample(channels=3)\n",
    "        self.heir_weight = nn.Parameter(torch.Tensor([1.0, 1.0, 1.0]))\n",
    "        self.debug = debug\n",
    "\n",
    "    def init_main_model(self):\n",
    "        args = self.args\n",
    "        filter_dim = args.filter_dim\n",
    "        latent_dim = args.filter_dim\n",
    "        im_size = args.im_size\n",
    "\n",
    "        self.conv1 = nn.Conv2d(3, filter_dim // 2, kernel_size=3, stride=1, padding=1)\n",
    "\n",
    "        self.res_1a = CondResBlock(args, filters=filter_dim // 2, latent_dim=latent_dim, im_size=im_size, downsample=True, classes=2, norm=args.norm, spec_norm=args.spec_norm)\n",
    "        self.res_1b = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, rescale=False, classes=2, norm=args.norm, spec_norm=args.spec_norm)\n",
    "\n",
    "        self.res_2a = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, downsample=True, rescale=False, classes=2, norm=args.norm, spec_norm=args.spec_norm)\n",
    "        self.res_2b = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, rescale=True, classes=2, norm=args.norm, spec_norm=args.spec_norm)\n",
    "\n",
    "        self.res_3a = CondResBlock(args, filters=2*filter_dim, latent_dim=latent_dim, im_size=im_size, downsample=False, classes=2, norm=args.norm, spec_norm=args.spec_norm)\n",
    "        self.res_3b = CondResBlock(args, filters=2*filter_dim, latent_dim=latent_dim, im_size=im_size, rescale=True, classes=2, norm=args.norm, spec_norm=args.spec_norm)\n",
    "\n",
    "        self.res_4a = CondResBlock(args, filters=4*filter_dim, latent_dim=latent_dim, im_size=im_size, downsample=False, classes=2, norm=args.norm, spec_norm=args.spec_norm)\n",
    "        self.res_4b = CondResBlock(args, filters=4*filter_dim, latent_dim=latent_dim, im_size=im_size, rescale=True, classes=2, norm=args.norm, spec_norm=args.spec_norm)\n",
    "\n",
    "        self.self_attn = Self_Attn(4 * filter_dim, self.act)\n",
    "\n",
    "        self.energy_map = nn.Linear(filter_dim*8, 1)\n",
    "\n",
    "    def init_mid_model(self):\n",
    "        args = self.args\n",
    "        filter_dim = args.filter_dim\n",
    "        latent_dim = args.filter_dim\n",
    "        im_size = args.im_size\n",
    "\n",
    "        self.mid_conv1 = nn.Conv2d(3, filter_dim, kernel_size=3, stride=1, padding=1)\n",
    "\n",
    "        self.mid_res_1a = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, downsample=True, rescale=False, classes=2)\n",
    "        self.mid_res_1b = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, rescale=False, classes=2)\n",
    "\n",
    "        self.mid_res_2a = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, downsample=True, rescale=False, classes=2)\n",
    "        self.mid_res_2b = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, rescale=True, classes=2)\n",
    "\n",
    "        self.mid_res_3a = CondResBlock(args, filters=2*filter_dim, latent_dim=latent_dim, im_size=im_size, downsample=False, classes=2)\n",
    "        self.mid_res_3b = CondResBlock(args, filters=2*filter_dim, latent_dim=latent_dim, im_size=im_size, rescale=True, classes=2)\n",
    "\n",
    "        self.mid_energy_map = nn.Linear(filter_dim*4, 1)\n",
    "        self.avg_pool = Downsample(channels=3)\n",
    "\n",
    "    def init_small_model(self):\n",
    "        args = self.args\n",
    "        filter_dim = args.filter_dim\n",
    "        latent_dim = args.filter_dim\n",
    "        im_size = args.im_size\n",
    "\n",
    "        self.small_conv1 = nn.Conv2d(3, filter_dim, kernel_size=3, stride=1, padding=1)\n",
    "\n",
    "        self.small_res_1a = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, downsample=True, rescale=False, classes=2)\n",
    "        self.small_res_1b = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, rescale=False, classes=2)\n",
    "\n",
    "        self.small_res_2a = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, downsample=True, rescale=False, classes=2)\n",
    "        self.small_res_2b = CondResBlock(args, filters=filter_dim, latent_dim=latent_dim, im_size=im_size, rescale=True, classes=2)\n",
    "\n",
    "        self.small_energy_map = nn.Linear(filter_dim*2, 1)\n",
    "\n",
    "    def main_model(self, x, latent):\n",
    "        x = self.act(self.conv1(x))\n",
    "\n",
    "        x = self.res_1a(x, latent)\n",
    "        x = self.res_1b(x, latent)\n",
    "\n",
    "        x = self.res_2a(x, latent)\n",
    "        x = self.res_2b(x, latent)\n",
    "\n",
    "\n",
    "        x = self.res_3a(x, latent)\n",
    "        x = self.res_3b(x, latent)\n",
    "\n",
    "        if self.args.self_attn:\n",
    "            x, _ = self.self_attn(x)\n",
    "\n",
    "        x = self.res_4a(x, latent)\n",
    "        x = self.res_4b(x, latent)\n",
    "        x = self.act(x)\n",
    "\n",
    "        x = x.mean(dim=2).mean(dim=2)\n",
    "\n",
    "        x = x.view(x.size(0), -1)\n",
    "        energy = self.energy_map(x)\n",
    "\n",
    "        if self.args.square_energy:\n",
    "            energy = torch.pow(energy, 2)\n",
    "\n",
    "        if self.args.sigmoid:\n",
    "            energy = F.sigmoid(energy)\n",
    "\n",
    "        return energy\n",
    "\n",
    "    def mid_model(self, x, latent):\n",
    "        x = F.avg_pool2d(x, 3, stride=2, padding=1)\n",
    "\n",
    "        x = self.act(self.mid_conv1(x))\n",
    "\n",
    "        x = self.mid_res_1a(x, latent)\n",
    "        x = self.mid_res_1b(x, latent)\n",
    "\n",
    "        x = self.mid_res_2a(x, latent)\n",
    "        x = self.mid_res_2b(x, latent)\n",
    "\n",
    "        x = self.mid_res_3a(x, latent)\n",
    "        x = self.mid_res_3b(x, latent)\n",
    "        x = self.act(x)\n",
    "\n",
    "        x = x.mean(dim=2).mean(dim=2)\n",
    "\n",
    "        x = x.view(x.size(0), -1)\n",
    "        energy = self.mid_energy_map(x)\n",
    "\n",
    "        if self.args.square_energy:\n",
    "            energy = torch.pow(energy, 2)\n",
    "\n",
    "        if self.args.sigmoid:\n",
    "            energy = F.sigmoid(energy)\n",
    "\n",
    "        return energy\n",
    "\n",
    "    def small_model(self, x, latent):\n",
    "        x = F.avg_pool2d(x, 3, stride=2, padding=1)\n",
    "        x = F.avg_pool2d(x, 3, stride=2, padding=1)\n",
    "\n",
    "        x = self.act(self.small_conv1(x))\n",
    "\n",
    "        x = self.small_res_1a(x, latent)\n",
    "        x = self.small_res_1b(x, latent)\n",
    "\n",
    "        x = self.small_res_2a(x, latent)\n",
    "        x = self.small_res_2b(x, latent)\n",
    "        x = self.act(x)\n",
    "\n",
    "        x = x.mean(dim=2).mean(dim=2)\n",
    "\n",
    "        x = x.view(x.size(0), -1)\n",
    "        energy = self.small_energy_map(x)\n",
    "\n",
    "        if self.args.square_energy:\n",
    "            energy = torch.pow(energy, 2)\n",
    "\n",
    "        if self.args.sigmoid:\n",
    "            energy = F.sigmoid(energy)\n",
    "\n",
    "        return energy\n",
    "\n",
    "    def label_map(self, latent):\n",
    "        x = self.act(self.map_fc1(latent))\n",
    "        x = self.act(self.map_fc2(x))\n",
    "        x = self.act(self.map_fc3(x))\n",
    "        x = self.act(self.map_fc4(x))\n",
    "\n",
    "        return x\n",
    "\n",
    "    def forward(self, x, latent):\n",
    "        args = self.args\n",
    "\n",
    "        if not self.cond:\n",
    "            latent = None\n",
    "\n",
    "        energy = self.main_model(x, latent)\n",
    "\n",
    "        if args.multiscale:\n",
    "            large_energy = energy\n",
    "            mid_energy = self.mid_model(x, latent)\n",
    "            small_energy = self.small_model(x, latent)\n",
    "            energy = torch.cat([small_energy, mid_energy, large_energy], dim=-1)\n",
    "\n",
    "        return energy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1.4 CEBM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CEBM(nn.Module):\n",
    "    \"\"\"\n",
    "    A generic class of CEBM. From \"Wu, Hao, et al. \"Conjugate Energy-Based Models.\" ICML 2021 \n",
    "    \"\"\"\n",
    "    def __init__(self, device, im_height, im_width, input_channels, channels, kernels, strides, paddings, hidden_dims, latent_dim, activation, **kwargs):\n",
    "        super().__init__()\n",
    "        self.device = device\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.conv_net = cnn_block(im_height, im_width, input_channels, channels, kernels, strides, paddings, activation, last_act=True, batchnorm=False, **kwargs)\n",
    "        out_h, out_w = cnn_output_shape(im_height, im_width, kernels, strides, paddings)\n",
    "        cnn_output_dim = out_h * out_w * channels[-1]\n",
    "        self.nss1_net = nn.Linear(cnn_output_dim, latent_dim)\n",
    "        self.nss2_net = nn.Linear(cnn_output_dim, latent_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.flatten(self.conv_net(x))\n",
    "        nss1 = self.nss1_net(h) \n",
    "        nss2 = self.nss2_net(h)\n",
    "        return nss1, -nss2**2\n",
    "\n",
    "    def log_partition(self, nat1, nat2):\n",
    "        \"\"\"\n",
    "        compute the log partition of a normal distribution\n",
    "        \"\"\"\n",
    "        return - 0.25 * (nat1 ** 2) / nat2 - 0.5 * (-2 * nat2).log()  \n",
    "\n",
    "    def nats_to_params(self, nat1, nat2):\n",
    "        \"\"\"\n",
    "        convert a Gaussian natural parameters its distritbuion parameters,\n",
    "        mu = - 0.5 *  (nat1 / nat2), \n",
    "        sigma = (- 0.5 / nat2).sqrt()\n",
    "        nat1 : natural parameter which correspond to x,\n",
    "        nat2 : natural parameter which correspond to x^2.      \n",
    "        \"\"\"\n",
    "        mu = - 0.5 * nat1 / nat2\n",
    "        sigma = (- 0.5 / nat2).sqrt()\n",
    "        return mu, sigma\n",
    "\n",
    "    def params_to_nats(self, mu, sigma):\n",
    "        \"\"\"\n",
    "        convert a Gaussian distribution parameters to the natrual parameters\n",
    "        nat1 = mean / sigma**2, \n",
    "        nat2 = - 1 / (2 * sigma**2)\n",
    "        nat1 : natural parameter which correspond to x,\n",
    "        nat2 : natural parameter which correspond to x^2.\n",
    "        \"\"\"\n",
    "        nat1 = mu / (sigma**2)\n",
    "        nat2 = - 0.5 / (sigma**2)\n",
    "        return nat1, nat2    \n",
    "\n",
    "    def log_factor(self, x, latents, expand_dim=None):\n",
    "        \"\"\"\n",
    "        compute the log factor log p(x | z) for the CEBM\n",
    "        \"\"\"\n",
    "        nss1, nss2 = self.forward(x)\n",
    "        if expand_dim is not None:\n",
    "            nss1 = nss1.expand(expand_dim , -1, -1)\n",
    "            nss2 = nss2.expand(expand_dim , -1, -1)\n",
    "            return (nss1 * latents).sum(2) + (nss2 * (latents**2)).sum(2)\n",
    "        else:\n",
    "            return (nss1 * latents).sum(1) + (nss2 * (latents**2)).sum(1) \n",
    "\n",
    "    def energy(self, x):\n",
    "        pass\n",
    "\n",
    "    def latent_params(self, x):\n",
    "        pass\n",
    "\n",
    "    def log_prior(self, latents):\n",
    "        pass   \n",
    "\n",
    "\n",
    "class CEBM_Gaussian(CEBM):\n",
    "    \"\"\"\n",
    "    conjugate EBM with a spherical Gaussian inductive bias\n",
    "    \"\"\"\n",
    "    def __init__(self, device, im_height, im_width, input_channels, channels, kernels, strides, paddings, hidden_dims, latent_dim, activation, **kwargs):\n",
    "        super().__init__(device, im_height, im_width, input_channels, channels, kernels, strides, paddings, hidden_dims, latent_dim, activation, **kwargs)\n",
    "        self.ib_mean = torch.zeros(latent_dim, device=self.device)\n",
    "        self.ib_log_std = torch.zeros(latent_dim, device=self.device)\n",
    "\n",
    "    def energy(self, x):\n",
    "        \"\"\"\n",
    "        return the energy of an input x\n",
    "        \"\"\"\n",
    "        nss1, nss2 = self.forward(x)\n",
    "        ib_nat1, ib_nat2 = self.params_to_nats(self.ib_mean, self.ib_log_std.exp())\n",
    "        logA_prior = self.log_partition(ib_nat1, ib_nat2)\n",
    "        logA_posterior = self.log_partition(ib_nat1+nss1, ib_nat2+nss2)\n",
    "        return logA_prior.sum(0) - logA_posterior.sum(1)   \n",
    "\n",
    "    def latent_params(self, x):\n",
    "        \"\"\"\n",
    "        return the posterior distribution parameters\n",
    "        \"\"\"\n",
    "        nss1, nss2 = self.forward(x)\n",
    "        ib_nat1, ib_nat2 = self.params_to_nats(self.ib_mean, self.ib_log_std.exp()) \n",
    "        return self.nats_to_params(ib_nat1+nss1, ib_nat2+nss2) \n",
    "\n",
    "class CEBM_GMM(CEBM):\n",
    "    \"\"\"\n",
    "    conjugate EBM with a GMM inductive bias\n",
    "    \"\"\"\n",
    "    def __init__(self, optimize_ib, num_clusters, device, im_height, im_width, input_channels, channels, kernels, strides, paddings, hidden_dims, latent_dim, activation, **kwargs):\n",
    "        super().__init__(device, im_height, im_width, input_channels, channels, kernels, strides, paddings, hidden_dims, latent_dim, activation, **kwargs)\n",
    "        #Suggested initialization\n",
    "        self.ib_means = torch.randn((num_clusters, latent_dim), device=self.device)\n",
    "        self.ib_log_stds = torch.ones((num_clusters, latent_dim), device=self.device).log()\n",
    "        if optimize_ib:\n",
    "            self.ib_means = nn.Parameter(self.ib_means)\n",
    "            self.ib_log_stds = nn.Parameter(self.ib_log_stds)\n",
    "        self.K = num_clusters\n",
    "        self.log_K = torch.tensor([self.K], device=self.device).log()\n",
    "\n",
    "    def energy(self, x):\n",
    "        \"\"\"\n",
    "        return the energy of an input x\n",
    "        \"\"\"\n",
    "        nss1, nss2 = self.forward(x)\n",
    "        ib_nat1, ib_nat2 = self.params_to_nats(self.ib_means, self.ib_log_stds.exp())\n",
    "        logA_prior = self.log_partition(ib_nat1, ib_nat2) # K * D\n",
    "        logA_posterior = self.log_partition(ib_nat1.unsqueeze(0)+nss1.unsqueeze(1), ib_nat2.unsqueeze(0)+nss2.unsqueeze(1)) # B * K * D\n",
    "        assert logA_prior.shape == (self.K, nss1.shape[1]), 'unexpected shape.'\n",
    "        assert logA_posterior.shape == (nss1.shape[0], self.K, nss1.shape[-1]), 'unexpected shape.'\n",
    "        return self.log_K - torch.logsumexp(logA_posterior.sum(2) - logA_prior.sum(1), dim=-1)   \n",
    "\n",
    "    def latent_params(self, x):\n",
    "        \"\"\"\n",
    "        return the posterior distribution parameters\n",
    "        \"\"\"\n",
    "        nss1, nss2 = self.forward(x)\n",
    "        ib_nat1, ib_nat2 = self.params_to_nats(self.ib_means, self.ib_log_stds.exp())\n",
    "        logA_prior = self.log_partition(ib_nat1, ib_nat2) # K * D\n",
    "        logA_posterior = self.log_partition(ib_nat1.unsqueeze(0)+nss1.unsqueeze(1), ib_nat2.unsqueeze(0)+nss2.unsqueeze(1)) # B * K * D\n",
    "        probs = torch.nn.functional.softmax(logA_posterior.sum(2) - logA_prior.sum(1), dim=-1)\n",
    "        means, stds = self.nats_to_params(ib_nat1.unsqueeze(0)+nss1.unsqueeze(1), ib_nat2.unsqueeze(0)+nss2.unsqueeze(1))\n",
    "        pred_y_expand = probs.argmax(-1).unsqueeze(-1).unsqueeze(-1).repeat(1, 1, means.shape[2])\n",
    "        return torch.gather(means, 1, pred_y_expand).squeeze(1), torch.gather(stds, 1, pred_y_expand).squeeze(1)\n",
    "\n",
    "\n",
    "def cnn_block(im_height, im_width, input_channels, channels, kernels, strides, paddings, activation, last_act, batchnorm, **kwargs):\n",
    "    \"\"\"\n",
    "    building blocks for a convnet.\n",
    "    each block is in form of:\n",
    "        Conv2d\n",
    "        BatchNorm2d(optinal)\n",
    "        Activation\n",
    "        Dropout(optional)\n",
    "    \"\"\"\n",
    "    if activation == 'Swish':\n",
    "        act = Swish()\n",
    "    elif activation == 'LeakyReLU':\n",
    "        act = nn.LeakyReLU(negative_slope=kwargs['leak_slope'], inplace=True)\n",
    "    else:\n",
    "        act = getattr(nn, activation)()\n",
    "    assert len(channels) == len(kernels), \"length of channels: %s,  length of kernels: %s\" % (len(channels), len(kernels))\n",
    "    assert len(channels) == len(strides), \"length of channels: %s,  length of strides: %s\" % (len(channels), len(strides))\n",
    "    assert len(channels) == len(paddings), \"length of channels: %s,  length of kernels: %s\" % (len(channels), len(paddings))\n",
    "    layers = []\n",
    "    in_c = input_channels\n",
    "    for i, out_c in enumerate(channels):\n",
    "        layers.append(nn.Conv2d(in_c, out_c, kernel_size=kernels[i], stride=strides[i], padding=paddings[i]))\n",
    "        if (i < (len(channels)-1)) or last_act:#Last layer will be customized \n",
    "            if batchnorm:\n",
    "                layers.append(nn.BatchNorm2d(out_c))\n",
    "            layers.append(act)\n",
    "            if 'dropout_prob' in kwargs:\n",
    "                layers.append(nn.Dropout2d(kwargs['dropout_prob']))\n",
    "            if 'maxpool_kernels' in kwargs and 'maxpool_strides' in kwargs:\n",
    "                layers.append(nn.MaxPool2d(kernel_size=kwargs['maxpool_kernels'][i], stride=kwargs['maxpool_strides'][i]))\n",
    "        in_c = out_c\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "def deconv_block(im_height, im_width, input_channels, channels, kernels, strides, paddings, activation, last_act, batchnorm, **kwargs):\n",
    "    \"\"\"\n",
    "    building blocks for a deconvnet\n",
    "    \"\"\"\n",
    "    if activation == 'Swish':\n",
    "        act = Swish()\n",
    "    elif activation == 'LeakyReLU':\n",
    "        act = nn.LeakyReLU(negative_slope=kwargs['leak_slope'], inplace=True)\n",
    "    else:\n",
    "        act = getattr(nn, activation)()\n",
    "    assert len(channels) == len(kernels), \"length of channels: %s,  length of kernels: %s\" % (len(channels), len(kernels))\n",
    "    assert len(channels) == len(strides), \"length of channels: %s,  length of strides: %s\" % (len(channels), len(strides))\n",
    "    assert len(channels) == len(paddings), \"length of channels: %s,  length of kernels: %s\" % (len(channels), len(paddings))\n",
    "    layers = []\n",
    "    in_c = input_channels\n",
    "    for i, out_c in enumerate(channels):\n",
    "        layers.append(nn.ConvTranspose2d(in_c, out_c, kernel_size=kernels[i], stride=strides[i], padding=paddings[i]))\n",
    "        if (i < (len(channels)-1)) or last_act:\n",
    "            if batchnorm:\n",
    "                layers.append(nn.BatchNorm2d(out_c)) \n",
    "            layers.append(act)\n",
    "        in_c = out_c\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "def mlp_block(input_dim, hidden_dims, activation, **kwargs):\n",
    "    \"\"\"\n",
    "    building blocks for a mlp\n",
    "    \"\"\"\n",
    "    if activation == 'Swish':\n",
    "        act = Swish()\n",
    "    elif activation == 'LeakyReLU':\n",
    "        act = nn.LeakyReLU(negative_slope=kwargs['leak_slope'], inplace=True)\n",
    "    else:\n",
    "        act = getattr(nn, activation)()\n",
    "    layers = []\n",
    "    in_dim = input_dim\n",
    "    for i, out_dim in enumerate(hidden_dims):\n",
    "        layers.append(nn.Linear(in_dim, out_dim))\n",
    "        layers.append(act)\n",
    "        in_dim = out_dim\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "class Identity(nn.Module):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__()\n",
    "    def forward(self, x):\n",
    "        return x\n",
    "\n",
    "class Reshape(nn.Module):\n",
    "    def __init__(self, shape):\n",
    "        super(Reshape, self).__init__()\n",
    "        self.shape = shape\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x.view(-1, *self.shape) \n",
    "\n",
    "def conv_output_shape(h_w, kernel_size=1, stride=1, padding=0, dilation=1):\n",
    "    \"\"\"\n",
    "    Utility function for computing output of convolutions\n",
    "    takes a tuple of (h,w) and returns a tuple of (h,w)\n",
    "    \"\"\"\n",
    "\n",
    "    if type(h_w) is not tuple:\n",
    "        h_w = (h_w, h_w)\n",
    "\n",
    "    if type(kernel_size) is not tuple:\n",
    "        kernel_size = (kernel_size, kernel_size)\n",
    "\n",
    "    if type(stride) is not tuple:\n",
    "        stride = (stride, stride)\n",
    "\n",
    "    if type(padding) is not tuple:\n",
    "        padding = (padding, padding)\n",
    "\n",
    "    h = (h_w[0] + (2 * padding[0]) - (dilation * (kernel_size[0] - 1)) - 1)// stride[0] + 1\n",
    "    w = (h_w[1] + (2 * padding[1]) - (dilation * (kernel_size[1] - 1)) - 1)// stride[1] + 1\n",
    "\n",
    "    return h, w\n",
    "\n",
    "def deconv_output_shape(h_w, kernel_size=1, stride=1, padding=0, dilation=1):\n",
    "    \"\"\"\n",
    "    Utility function for computing output of deconvolutions\n",
    "    takes a tuple of (h,w) and returns a tuple of (h,w)\n",
    "    \"\"\"\n",
    "\n",
    "    if type(h_w) is not tuple:\n",
    "        h_w = (h_w, h_w)\n",
    "\n",
    "    if type(kernel_size) is not tuple:\n",
    "        kernel_size = (kernel_size, kernel_size)\n",
    "\n",
    "    if type(stride) is not tuple:\n",
    "        stride = (stride, stride)\n",
    "\n",
    "    if type(padding) is not tuple:\n",
    "        padding = (padding, padding)\n",
    "    h = (h_w[0] - 1) * stride[0] - 2 * padding[0]  + (dilation * (kernel_size[0] - 1)) + 1\n",
    "    w = (h_w[1] - 1) * stride[1] - 2 * padding[1]  + (dilation * (kernel_size[1] - 1)) + 1\n",
    "\n",
    "    return h, w\n",
    "\n",
    "def cnn_output_shape(h, w, kernels, strides, paddings):\n",
    "    h_w = (h, w)\n",
    "    for i, kernel in enumerate(kernels):\n",
    "        h_w = conv_output_shape(h_w, kernels[i], strides[i], paddings[i])\n",
    "    return h_w\n",
    "\n",
    "def dcnn_output_shape(h, w, kernels, strides, paddings):\n",
    "    h_w = (h, w)\n",
    "    for i, kernel in enumerate(kernels):\n",
    "        h_w = deconv_output_shape(h_w, kernels[i], strides[i], paddings[i])\n",
    "    return h_w"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1.5 GraphEBM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphEBM(nn.Module):\n",
    "    \"\"\"\n",
    "    E.g. GraphEBM(models_dict), where\n",
    "        models = {\n",
    "            \"re_0\": model_re.set_c(OPERATORS[\"SameShape\"].get_node_repr()[None]),\n",
    "            \"re_1\": model_re.set_c(OPERATORS[\"IsInside\"].get_node_repr()[None]),\n",
    "            \"obj_0\": model_concept.set_c(CONCEPTS[\"Rect\"].get_node_repr()[None]),\n",
    "        }\n",
    "        assign_dict = {\n",
    "            \"obj_0\": {2},\n",
    "            \"re_0\": {(0, 1)},\n",
    "            \"re_1\": {(1, 2)},\n",
    "        }\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        models,\n",
    "        assign_dict,\n",
    "        mask_arity,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.models = models\n",
    "        self.assign_dict = assign_dict\n",
    "        self.mask_arity = mask_arity\n",
    "\n",
    "    def to(self, device):\n",
    "        self.device = device\n",
    "        super().to(device)\n",
    "        return self\n",
    "\n",
    "    def init_assign_dict_param(self):\n",
    "        self.assign_dict_param = nn.ParameterDict()\n",
    "        for key, List in self.assign_dict.items():\n",
    "            self.assign_dict_param[\"{}^{}\".format(key, List)] = nn.Parameter(1+torch.randn(1).to(self.device)*0.01)\n",
    "\n",
    "    def del_assign_dict_param(self):\n",
    "        delattr(self, \"assign_dict_param\")\n",
    "\n",
    "    def forward(self, input, mask, c_repr=None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            input:  x, [B, C, H, W]\n",
    "            mask:   a list of masks, each with shape [B, 1, H, W]\n",
    "            c_repr: embedding, [B, REPR_DIM]\n",
    "        \"\"\"\n",
    "        energy = 0\n",
    "        self.info = {}\n",
    "        for key, mask_ids in self.assign_dict.items():\n",
    "            assert isinstance(mask_ids, set)\n",
    "            for mask_ids_setele in mask_ids:\n",
    "                if not isinstance(mask_ids_setele, tuple):\n",
    "                    mask_ids_setele = (mask_ids_setele,)\n",
    "                mask_ele = tuple(mask[id] for id in mask_ids_setele)\n",
    "                if len(mask_ele) == 1:\n",
    "                    input_ele = input\n",
    "                else:\n",
    "                    assert len(mask_ele) == 2\n",
    "                    input_ele = (input, input)\n",
    "                energy_ele = self.models[key](input_ele, mask_ele)\n",
    "                if hasattr(self, \"assign_dict_param\"):\n",
    "                    energy_ele = energy_ele * self.assign_dict_param[\"{}^{}\".format(key, mask_ids)]\n",
    "                energy = energy + energy_ele\n",
    "                string = \",\".join([str(ele) for ele in mask_ids_setele])\n",
    "                string = \"({})\".format(string) if len(mask_ids_setele) > 1 else string\n",
    "                self.info[\"{}^{}\".format(key, string)] = to_np_array(energy_ele)\n",
    "        return energy\n",
    "\n",
    "    def infer_recur(\n",
    "        self,\n",
    "        model,\n",
    "        input,\n",
    "        mask,\n",
    "        max_recur_depth=6,\n",
    "    ):\n",
    "        \"\"\"Recursively infer the concept component of the mask.\"\"\"\n",
    "        mask_list = []\n",
    "        for i in range(max_recur_depth):\n",
    "            neg_mask_ensemble_sorted, neg_out_ensemble_sorted = model.ground(input, args=args_concept, mask=mask, ensemble_size=1)\n",
    "            mask_c = tuple(mask_ele[:,0] for mask_ele in neg_mask_ensemble_sorted)\n",
    "            mask_c = and_mask(mask_c, mask)\n",
    "            mask = subtract_mask(mask, mask_c)\n",
    "            mask_list.append(mask_c)\n",
    "            if mask[0].sum() == 0:\n",
    "                print(\"break at {}\".format(i))\n",
    "                break\n",
    "        masks_c = Zip(*mask_list, function=torch.cat)\n",
    "        return masks_c\n",
    "\n",
    "    \n",
    "    def classify(\n",
    "        self,\n",
    "        input,\n",
    "        mask,\n",
    "        concept_collection,\n",
    "        CONCEPTS=None,\n",
    "        OPERATORS=None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Uses a corresponding model to classify each mask in self.assign_dict.\n",
    "        \"\"\"\n",
    "        \n",
    "        pred_dict = {}\n",
    "        for key, mask_ids in self.assign_dict.items():\n",
    "            assert isinstance(mask_ids, set)\n",
    "            for mask_ids_setele in mask_ids:\n",
    "                if not isinstance(mask_ids_setele, tuple):\n",
    "                    mask_ids_setele = (mask_ids_setele,)\n",
    "                mask_ele = tuple(mask[id] for id in mask_ids_setele)\n",
    "                if len(mask_ele) == 1:\n",
    "                    input_ele = input\n",
    "                else:\n",
    "                    assert len(mask_ele) == 2\n",
    "                    input_ele = (input, input)\n",
    "                pred_ele = self.models[key].classify(input_ele, mask_ele, concept_collection=concept_collection[key], CONCEPTS=CONCEPTS, OPERATORS=OPERATORS)\n",
    "                string = \",\".join([str(ele) for ele in mask_ids_setele])\n",
    "                string = \"({})\".format(string) if len(mask_ids_setele) > 1 else string\n",
    "                pred_dict[\"{}^{}\".format(key, string)] = pred_ele\n",
    "        return pred_dict\n",
    "\n",
    "    def ground(\n",
    "        self,\n",
    "        input,\n",
    "        args,\n",
    "        mask=None,\n",
    "        c_repr=None,\n",
    "        z=None,\n",
    "        ensemble_size=18,\n",
    "        topk=-1,\n",
    "        w_init_type=\"random\",\n",
    "        sample_step=150,\n",
    "        isplot=2,\n",
    "        ground_truth_mask: tuple = None,\n",
    "        **kwargs\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Given input (and optionally initial masks), find the best masks that\n",
    "        minimize the total energy.\n",
    "\n",
    "        Args:\n",
    "            isplot: whether or not to plot. If 0, skips plotting. If 1,\n",
    "                only plots individual discovered masks. If 2, plots learning curve as well.\n",
    "\n",
    "            ground_truth_mask: the real (positive) mask. If given, ground()\n",
    "                will plot the energy of the ground truth masks as well. Tuple of tensors\n",
    "                with shape [1, C, H, W] and length mask_arity.\n",
    "        \"\"\"\n",
    "        def init_neg_mask(input, init, ensemble_size):\n",
    "            \"\"\"Initialize negative mask\"\"\"\n",
    "            device = input.device\n",
    "            neg_mask = tuple(torch.rand(input.shape[0]*ensemble_size, 1, *input.shape[2:]).to(device) for k in range(self.mask_arity))\n",
    "            if init == \"input-mask\":\n",
    "                assert input.shape[1] == 10\n",
    "                input_l = repeat_n(input, n_repeats=ensemble_size)\n",
    "                neg_mask = tuple(neg_mask[k] * (input_l.argmax(1)[:, None] != 0) for k in range(self.mask_arity))\n",
    "            for k in range(self.mask_arity):\n",
    "                neg_mask[k].requires_grad = True\n",
    "            return neg_mask\n",
    "\n",
    "        def plot_discovered_mask_summary(num_examples: int, neg_mask_ensemble_sorted: torch.Tensor, should_quantize: bool):\n",
    "            plt.figure(figsize=(18,3))\n",
    "            for batch_idx in range(len(neg_mask_ensemble_sorted[0])):\n",
    "                for ex in range(num_examples):\n",
    "                    for mask_idx in range(len(neg_mask_ensemble_sorted)):\n",
    "                        ax = plt.subplot(1, num_examples, ex + 1)\n",
    "                        # Pull single-channel (0)\n",
    "                        mask = neg_mask_ensemble_sorted[mask_idx][batch_idx][ex][0].cpu()\n",
    "                        image = np.zeros((*mask.shape, 4)) # (H, W, C, alpha)\n",
    "                        color = np.asarray(matplotlib.colors.to_rgb(COLOR_LIST[mask_idx]))\n",
    "                        for h in range(mask.shape[0]):\n",
    "                            for w in range(mask.shape[1]):\n",
    "                                opacity = torch.round(mask[h][w]) if should_quantize else mask[h][w]\n",
    "                                pixel = opacity * np.asarray((*color, 0.5)) # add alpha channel\n",
    "                                image[h][w] = pixel\n",
    "                        plt.imshow(image)\n",
    "                        ax.set_title(\"E: {:.5f}\\n\".format(neg_out_ensemble_sorted[batch_idx][ex]))\n",
    "            plt.show()\n",
    "\n",
    "        assert (not isinstance(input, tuple)) and (not isinstance(input, list)) and len(input.shape) == 4\n",
    "        device = input.device\n",
    "        # Update args:\n",
    "        args = deepcopy(args)\n",
    "        for key, value in kwargs.items():\n",
    "            setattr(args, key, value)\n",
    "        args.sample_step = sample_step\n",
    "        z_mode = self.models[next(iter(self.models))].z_mode\n",
    "        args.ebm_target = \"mask\" if z_mode == \"None\" else \"mask+z\"\n",
    "\n",
    "        if ground_truth_mask is not None:\n",
    "            print(\"Ground truth mask energies:\")\n",
    "            energy = self.forward(input, ground_truth_mask, c_repr)\n",
    "            # Plot all masks together, followed by each individual mask\n",
    "            visualize_matrices(torch.cat([sum(ground_truth_mask), *ground_truth_mask], dim=0).squeeze(1),\n",
    "                               images_per_row=len(ground_truth_mask) + 1,\n",
    "                               subtitles=[\n",
    "                                   # Print overall energy\n",
    "                                   \"E: {:.5f}\\n\".format(float(energy)) + \\\n",
    "                                    # Print energy of individual components\n",
    "                                    \"\\n\".join([\"{}: {:.5f}\".format(key, float(self.info[key])) for key in self.info])] + [\"\"] * len(ground_truth_mask))\n",
    "\n",
    "        print(\"Inferred masks:\")\n",
    "        # Perform SGLD:\n",
    "        if mask is None:\n",
    "            neg_mask = init_neg_mask(input, init=w_init_type, ensemble_size=ensemble_size)\n",
    "        else:\n",
    "            neg_mask = mask\n",
    "        (img_ensemble, neg_mask_ensemble, z_ensemble), neg_out_list_ensemble = neg_mask_sgd_ensemble(\n",
    "            self, input, neg_mask, c_repr, z=z, args=args,\n",
    "            ensemble_size=ensemble_size,\n",
    "        )\n",
    "        neg_out_ensemble = neg_out_list_ensemble[-1]  # [time_steps, ensemble_size, B]\n",
    "        # Sort the obtained results by energy for each example:\n",
    "        neg_out_ensemble = torch.FloatTensor(neg_out_ensemble).transpose(0,1)\n",
    "        neg_out_argsort = neg_out_ensemble.argsort(1)  # [B, ensemble_size]\n",
    "        batch_size = neg_out_argsort.shape[0]\n",
    "        neg_out_ensemble_sorted = torch.stack([neg_out_ensemble[i][neg_out_argsort[i]] for i in range(len(neg_mask_ensemble[0]))])  # [B, ensemble_size]\n",
    "\n",
    "        if img_ensemble is not None:\n",
    "            if args.is_image_tuple:\n",
    "                img_ensemble = tuple(img_ensemble[k].transpose(0,1) for k in range(len(img)))  # Each element [B, ensemble_size, C, H, W]\n",
    "                img_ensemble_sorted = []\n",
    "                for k in range(len(img)):\n",
    "                    img_ensemble_sorted.append(torch.stack([img_ensemble[k][i][neg_out_argsort[i]] for i in range(batch_size)]))\n",
    "                img_ensemble_sorted = tuple(img_ensemble_sorted)  # each element: [B, ensemble_size, C, H, W] sorted along dim=1 according to neg_out\n",
    "            else:\n",
    "                img_ensemble = img_ensemble.transpose(0,1)  # [B, ensemble_size, C, H, W]\n",
    "                img_ensemble_sorted = torch.stack([img_ensemble[i][neg_out_argsort[i]] for i in range(batch_size)])\n",
    "        else:\n",
    "            img_ensemble_sorted = None\n",
    "\n",
    "        if neg_mask_ensemble is not None:\n",
    "            neg_mask_ensemble = tuple(neg_mask_ensemble[k].transpose(0,1) for k in range(self.mask_arity))  # Each element [B, ensemble_size, C, H, W]\n",
    "            neg_mask_ensemble_sorted = []\n",
    "            for k in range(self.mask_arity):\n",
    "                neg_mask_ensemble_sorted.append(torch.stack([neg_mask_ensemble[k][i][neg_out_argsort[i]] for i in range(len(neg_mask_ensemble[0]))]))\n",
    "            # Shape: tuple of length mask_arity\n",
    "            # where each element: [B, ensemble_size, C, H, W] sorted along dim=1 according to neg_out\n",
    "            # (aka, lowest energy configuration first)\n",
    "            neg_mask_ensemble_sorted = tuple(neg_mask_ensemble_sorted)  # each element: [B, ensemble_size, C, H, W] sorted along dim=1 according to neg_out\n",
    "        else:\n",
    "            neg_mask_ensemble_sorted = None\n",
    "\n",
    "        if z_ensemble is not None:\n",
    "            z_ensemble = tuple(z_ensemble[k].transpose(0,1) for k in range(len(z_ensemble)))  # Each element [B, ensemble_size, Z]\n",
    "            z_ensemble_sorted = []\n",
    "            for k in range(len(z_ensemble)):\n",
    "                z_ensemble_sorted.append(torch.stack([z_ensemble[k][i][neg_out_argsort[i]] for i in range(batch_size)]))\n",
    "            z_ensemble_sorted = tuple(z_ensemble_sorted)  # each element: [B, ensemble_size, Z] sorted along dim=1 according to neg_out\n",
    "        else:\n",
    "            z_ensemble_sorted = None\n",
    "\n",
    "        # Obtain each individual energy for component models:\n",
    "        info = {}\n",
    "        neg_out_argsort = to_np_array(neg_out_argsort)\n",
    "        for key, value in self.info.items():\n",
    "            value_reshape = value.reshape(ensemble_size, -1).T  # [B, ensemble_size]\n",
    "            info[key] = []\n",
    "            for i in range(len(value_reshape)):\n",
    "                info[key].append(value_reshape[i][neg_out_argsort[i]])\n",
    "            info[key] = np.stack(info[key])\n",
    "\n",
    "        NUM_PREV_EXAMPLES = 6\n",
    "\n",
    "        if isplot >= 2:\n",
    "            # Plot SGLD learning curve:\n",
    "            plt.figure(figsize=(12,6))\n",
    "            for i in range(min(neg_out_list_ensemble.shape[-1], 6)):  # neg_out_list_ensemble: [sample_step, ensemble_size, B]\n",
    "                print(\"Example {}\".format(i))\n",
    "                for k in range(6):\n",
    "                    plt.plot(neg_out_list_ensemble[:,neg_out_argsort[i][k],i], c=COLOR_LIST[k], label=\"id_{}\".format(k), alpha=0.4)\n",
    "            plt.legend()\n",
    "            plt.show()\n",
    "        if isplot >= 1:\n",
    "            # Show original input image for reference\n",
    "            print(\"Original inputs:\")\n",
    "            visualize_matrices(input.argmax(1).repeat_interleave(NUM_PREV_EXAMPLES, 0))\n",
    "\n",
    "            # Plot a summary plot, superimposing different masks such that each mask has a different color\n",
    "            print(f\"Top-{NUM_PREV_EXAMPLES} lowest-energy mask sets, all plotted together\")\n",
    "            print(f\"Key parameters: SGLD_mutual_exclusive_coef={str(args.SGLD_mutual_exclusive_coef)}\",\n",
    "                    f\"SGLD_object_exceed_coef={str(args.SGLD_object_exceed_coef)}\")\n",
    "            plot_discovered_mask_summary(NUM_PREV_EXAMPLES, neg_mask_ensemble_sorted, should_quantize=False)\n",
    "\n",
    "            # Plot the same plot, but this time quantized so there are no color gradations\n",
    "            print(\"Quantized plot:\")\n",
    "            plot_discovered_mask_summary(NUM_PREV_EXAMPLES, neg_mask_ensemble_sorted, should_quantize=True)\n",
    "\n",
    "            # For each batch element\n",
    "            for i in range(len(neg_mask_ensemble_sorted[0])):\n",
    "                print(\"Example {}\".format(i))\n",
    "                # Loop through each mask (show them horizontally)\n",
    "                for k in range(len(neg_mask_ensemble_sorted)):\n",
    "                    visualize_matrices(\n",
    "                        torch.round(neg_mask_ensemble_sorted[k][i][:NUM_PREV_EXAMPLES].squeeze(1)), images_per_row=NUM_PREV_EXAMPLES,\n",
    "                        subtitles=[\"E: {:.5f}\\n\".format(neg_out_ensemble_sorted[i][j]) + \"\\n\".join([\"{}: {:.5f}\".format(key, info[key][i][j]) for key in info]) for j in range(NUM_PREV_EXAMPLES)] if k == 0 else None\n",
    "                    )\n",
    "        return (img_ensemble_sorted, neg_mask_ensemble_sorted, z_ensemble_sorted), neg_out_ensemble_sorted, info\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": self.__class__.__name__}\n",
    "        model_dict[\"models\"] = {key: model.model_dict for key, model in self.models.items()}\n",
    "        model_dict[\"assign_dict\"] = self.assign_dict\n",
    "        model_dict[\"mask_arity\"] = self.mask_arity\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "\n",
    "class SumEBM(nn.Module):\n",
    "    \"\"\"\n",
    "    SumEBM((model1, c_repr1, \"Line\"), (model2, c_repr2, \"Vertical\"))\n",
    "    \"\"\"\n",
    "    def __init__(self, *models):\n",
    "        super().__init__()\n",
    "        self.models = nn.ModuleList(models)\n",
    "        self.modes = [model.mode for model in self.models]\n",
    "        self.is_two_branch = True if \"relation\" in self.modes or \"operator\" in self.modes else False\n",
    "\n",
    "    def forward(self, input, mask, c_repr):\n",
    "        energy_list = []\n",
    "        for model in self.models:\n",
    "            if self.is_two_branch:\n",
    "                if model.mode in [\"relation\", \"operator\"]:\n",
    "                    energy = model(input, mask, model.c_repr)\n",
    "                elif model.mode == \"concept\":\n",
    "                    energy = model(input[0], (mask[0],), model.c_repr) + model(input[1], (mask[1],), model.c_repr)\n",
    "                else:\n",
    "                    raise Exception(\"The model's mode must be one of 'concept', 'relation' or 'operator'.\")\n",
    "            else:\n",
    "                energy = model(input, mask, model.c_repr)\n",
    "            energy_list.append(energy)\n",
    "        energy = torch.cat(energy_list, -1).sum(-1, keepdims=True)\n",
    "        return energy\n",
    "\n",
    "    def __add__(self, model):\n",
    "        if model.__class__.__name__ == \"ConceptEBM\":\n",
    "            self.models.append(model)\n",
    "            self.modes.append(model.mode)\n",
    "        elif model.__class__.__name__ == \"SumEBM\":\n",
    "            self.models = self.models + model.models\n",
    "            self.modes += model.modes\n",
    "        else:\n",
    "            raise Exception(\"model is not valid!\")\n",
    "        self.is_two_branch = True if \"relation\" in self.modes or \"operator\" in self.modes else False\n",
    "        return self\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": self.__class__.__name__}\n",
    "        model_dict[\"models\"] = [model.model_dict for model in self.models]\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "def and_mask(mask, mask_ref):\n",
    "    \"\"\"Perform And operation that only preserve the non-zero part of mask_ref onto mask.\"\"\"\n",
    "    mask_and = []\n",
    "    for mask1_ele, mask2_ele in zip(mask, mask_ref):\n",
    "        mask_and_ele = (mask1_ele * mask2_ele.round().bool()).float()\n",
    "        mask_and.append(mask_and_ele)\n",
    "    return tuple(mask_and)\n",
    "\n",
    "def subtract_mask(mask1, mask2):\n",
    "    \"\"\"Subtract mask1 by mask2 as Boolean.\"\"\"\n",
    "    mask_sub = []\n",
    "    for mask1_ele, mask2_ele in zip(mask1, mask2):\n",
    "        mask_sub_ele = (mask1_ele.round().bool() & ~mask2_ele.round().bool()).float()\n",
    "        mask_sub.append(mask_sub_ele)\n",
    "    return tuple(mask_sub)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 Modules:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2.1 SpectralNorm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SpectralNorm(object):\n",
    "    def __init__(self, name, bound=False):\n",
    "        self.name = name\n",
    "        self.bound = bound\n",
    "\n",
    "    def compute_weight(self, module):\n",
    "        weight = getattr(module, self.name + '_orig')\n",
    "        u = getattr(module, self.name + '_u')\n",
    "        size = weight.size()\n",
    "        weight_mat = weight.contiguous().view(size[0], -1)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            v = weight_mat.t() @ u\n",
    "            v = v / v.norm()\n",
    "            u = weight_mat @ v\n",
    "            u = u / u.norm()\n",
    "\n",
    "        sigma = u @ weight_mat @ v\n",
    "\n",
    "        if self.bound:\n",
    "            weight_sn = weight / (sigma + 1e-6) * torch.clamp(sigma, max=1)\n",
    "\n",
    "        else:\n",
    "            weight_sn = weight / sigma\n",
    "\n",
    "        return weight_sn, u\n",
    "\n",
    "    @staticmethod\n",
    "    def apply(module, name, bound):\n",
    "        fn = SpectralNorm(name, bound)\n",
    "\n",
    "        weight = getattr(module, name)\n",
    "        del module._parameters[name]\n",
    "        module.register_parameter(name + '_orig', weight)\n",
    "        input_size = weight.size(0)\n",
    "        u = weight.new_empty(input_size).normal_()\n",
    "        module.register_buffer(name, weight)\n",
    "        module.register_buffer(name + '_u', u)\n",
    "\n",
    "        module.register_forward_pre_hook(fn)\n",
    "\n",
    "        return fn\n",
    "\n",
    "    def __call__(self, module, input):\n",
    "        weight_sn, u = self.compute_weight(module)\n",
    "        setattr(module, self.name, weight_sn)\n",
    "        setattr(module, self.name + '_u', u)\n",
    "\n",
    "\n",
    "def spectral_norm(module, init=True, std=1, bound=False):\n",
    "    if init:\n",
    "        nn.init.normal_(module.weight, 0, std)\n",
    "\n",
    "    if hasattr(module, 'bias') and module.bias is not None:\n",
    "        module.bias.data.zero_()\n",
    "\n",
    "    SpectralNorm.apply(module, 'weight', bound=bound)\n",
    "\n",
    "    return module\n",
    "\n",
    "\n",
    "class WSConv2d(nn.Conv2d):\n",
    "    \"\"\"Conv layer with the output normalized.\"\"\"\n",
    "    def __init__(self, in_channels, out_channels, kernel_size, stride=1,\n",
    "                 padding=0, dilation=1, groups=1, bias=True):\n",
    "        super(WSConv2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)\n",
    "\n",
    "    def forward(self, x):\n",
    "        weight = self.weight\n",
    "        weight_mean = weight.mean((1,2,3), keepdims=True)\n",
    "        weight = weight - weight_mean\n",
    "        std = weight.view(weight.size(0), -1).std(dim=1).view(-1, 1, 1, 1) + 1e-5\n",
    "        weight = weight / std.expand_as(weight)\n",
    "        return F.conv2d(x, weight, self.bias, self.stride,\n",
    "                        self.padding, self.dilation, self.groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2.2 ResBlock:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, in_channel, out_channel, n_class=None, downsample=False, is_spec_norm=True):\n",
    "        super().__init__()\n",
    "\n",
    "        if is_spec_norm:\n",
    "            self.conv1 = spectral_norm(\n",
    "                nn.Conv2d(\n",
    "                    in_channel,\n",
    "                    out_channel,\n",
    "                    3,\n",
    "                    padding=1,\n",
    "                    bias=False if n_class is not None else True,\n",
    "                )\n",
    "            )\n",
    "            self.conv2 = spectral_norm(\n",
    "                nn.Conv2d(\n",
    "                    out_channel,\n",
    "                    out_channel,\n",
    "                    3,\n",
    "                    padding=1,\n",
    "                    bias=False if n_class is not None else True,\n",
    "                ), std=1e-10, bound=True\n",
    "            )\n",
    "        else:\n",
    "            self.conv1 = nn.Conv2d(\n",
    "                in_channel,\n",
    "                out_channel,\n",
    "                3,\n",
    "                padding=1,\n",
    "                bias=False if n_class is not None else True,\n",
    "            )\n",
    "            self.conv2 = nn.Conv2d(\n",
    "                out_channel,\n",
    "                out_channel,\n",
    "                3,\n",
    "                padding=1,\n",
    "                bias=False if n_class is not None else True,\n",
    "            )\n",
    "\n",
    "        self.class_embed = None\n",
    "\n",
    "        if n_class is not None:\n",
    "            class_embed = nn.Embedding(n_class, out_channel * 2 * 2)\n",
    "            class_embed.weight.data[:, : out_channel * 2] = 1\n",
    "            class_embed.weight.data[:, out_channel * 2 :] = 0\n",
    "\n",
    "            self.class_embed = class_embed\n",
    "\n",
    "        self.skip = None\n",
    "\n",
    "        if in_channel != out_channel or downsample:\n",
    "            if is_spec_norm:\n",
    "                self.skip = nn.Sequential(\n",
    "                    spectral_norm(nn.Conv2d(in_channel, out_channel, 1, bias=False))\n",
    "                )\n",
    "            else:\n",
    "                self.skip = nn.Sequential(\n",
    "                    nn.Conv2d(in_channel, out_channel, 1, bias=False)\n",
    "                )\n",
    "\n",
    "        self.downsample = downsample\n",
    "\n",
    "    def forward(self, input, class_id=None):\n",
    "        out = input\n",
    "\n",
    "        out = self.conv1(out)\n",
    "\n",
    "        if self.class_embed is not None:\n",
    "            embed = self.class_embed(class_id).view(input.shape[0], -1, 1, 1)\n",
    "            weight1, weight2, bias1, bias2 = embed.chunk(4, 1)\n",
    "            out = weight1 * out + bias1\n",
    "\n",
    "        out = F.leaky_relu(out, negative_slope=0.2)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "\n",
    "        if self.class_embed is not None:\n",
    "            out = weight2 * out + bias2\n",
    "\n",
    "        if self.skip is not None:\n",
    "            skip = self.skip(input)\n",
    "\n",
    "        else:\n",
    "            skip = input\n",
    "\n",
    "        out = out + skip\n",
    "\n",
    "        if self.downsample:\n",
    "            out = F.avg_pool2d(out, 2)\n",
    "\n",
    "        out = F.leaky_relu(out, negative_slope=0.2)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2.3 CResBlock:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CResBlock(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channel,\n",
    "        out_channel,\n",
    "        repr_dim=None,\n",
    "        downsample=False,\n",
    "        is_spec_norm=True,\n",
    "        c_repr_mode=\"l1\",\n",
    "        c_repr_base=2,\n",
    "        z_mode=\"None\",\n",
    "        z_dim=4,\n",
    "        img_dims=2,\n",
    "        act_name=\"leakyrelu0.2\",\n",
    "        normalization_type=\"None\",\n",
    "        dropout=0,\n",
    "        is_res=True,\n",
    "        downsample_mode=\"None\",\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        in_channel_combined = in_channel\n",
    "        out_channel_combined = out_channel\n",
    "        if c_repr_mode.startswith(\"c\"):\n",
    "            in_channel_combined += c_repr_base\n",
    "            out_channel_combined += c_repr_base\n",
    "        if z_mode.startswith(\"c\"):\n",
    "            in_channel_combined += z_dim\n",
    "            out_channel_combined += z_dim\n",
    "        self.img_dims = img_dims\n",
    "        if img_dims == 2:\n",
    "            kernel_size = 3\n",
    "            padding = 1\n",
    "            self.pool_kernel_size = 2\n",
    "        elif img_dims == 1:\n",
    "            kernel_size = (3, 1)\n",
    "            padding = (1, 0)\n",
    "            self.pool_kernel_size = (2, 1)\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        self.act1 = get_activation(act_name)\n",
    "        self.act2 = get_activation(act_name)\n",
    "        self.bn1 = get_normalization(normalization_type, in_channels=out_channel)\n",
    "        self.bn2 = get_normalization(normalization_type, in_channels=out_channel)\n",
    "        self.dropout = dropout\n",
    "        self.is_res = is_res\n",
    "        self.downsample = downsample\n",
    "        self.downsample_mode = downsample_mode\n",
    "        if self.dropout != 0:\n",
    "            self.dropout_fn = nn.Dropout(dropout)\n",
    "\n",
    "        if is_spec_norm in [True, \"True\"]:\n",
    "            self.conv1 = spectral_norm(\n",
    "                nn.Conv2d(\n",
    "                    in_channel_combined,\n",
    "                    out_channel,\n",
    "                    kernel_size,\n",
    "                    padding=padding,\n",
    "                    bias=False,\n",
    "                )\n",
    "            )\n",
    "            self.conv2 = spectral_norm(\n",
    "                nn.Conv2d(\n",
    "                    out_channel_combined,\n",
    "                    out_channel,\n",
    "                    kernel_size,\n",
    "                    padding=padding,\n",
    "                    bias=False,\n",
    "                ), std=1e-10, bound=True\n",
    "            )\n",
    "        elif is_spec_norm in [False, \"False\"]:\n",
    "            self.conv1 = nn.Conv2d(\n",
    "                in_channel_combined,\n",
    "                out_channel,\n",
    "                kernel_size,\n",
    "                padding=padding,\n",
    "                bias=False,\n",
    "            )\n",
    "            self.conv2 = nn.Conv2d(\n",
    "                out_channel_combined,\n",
    "                out_channel,\n",
    "                kernel_size,\n",
    "                padding=padding,\n",
    "                bias=False,\n",
    "            )\n",
    "        elif is_spec_norm == \"ws\":\n",
    "            self.conv1 = WSConv2d(\n",
    "                in_channel_combined,\n",
    "                out_channel,\n",
    "                kernel_size,\n",
    "                padding=padding,\n",
    "                bias=False,\n",
    "            )\n",
    "            self.conv2 = WSConv2d(\n",
    "                out_channel_combined,\n",
    "                out_channel,\n",
    "                kernel_size,\n",
    "                padding=padding,\n",
    "                bias=False,\n",
    "            )\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        # Constructing self.class_embed:\n",
    "        self.c_repr_mode = c_repr_mode\n",
    "        self.c_repr_base = c_repr_base\n",
    "        if c_repr_mode.startswith(\"l1\"):\n",
    "            self.class_embed = nn.Linear(repr_dim, out_channel * 4)\n",
    "        elif c_repr_mode.startswith(\"l2\"):\n",
    "            self.class_embed = nn.Sequential(nn.Linear(repr_dim, repr_dim * c_repr_base*2),\n",
    "                                             nn.LeakyReLU(negative_slope=0.2),\n",
    "                                             nn.Linear(repr_dim * c_repr_base*2, out_channel * 4),\n",
    "                                            )\n",
    "        elif c_repr_mode.startswith(\"l3\"):\n",
    "            self.class_embed = nn.Sequential(nn.Linear(repr_dim, repr_dim * c_repr_base*2),\n",
    "                                             nn.LeakyReLU(negative_slope=0.2),\n",
    "                                             nn.Linear(repr_dim * c_repr_base*2, repr_dim * c_repr_base*2),\n",
    "                                             nn.LeakyReLU(negative_slope=0.2),\n",
    "                                             nn.Linear(repr_dim * c_repr_base*2, out_channel * 4),\n",
    "                                            )\n",
    "        elif c_repr_mode.startswith(\"c1\"):\n",
    "            self.class_embed = nn.Linear(repr_dim, c_repr_base*2)\n",
    "        elif c_repr_mode.startswith(\"c2\"):\n",
    "            self.class_embed = nn.Sequential(nn.Linear(repr_dim, repr_dim * c_repr_base*2),\n",
    "                                             nn.LeakyReLU(negative_slope=0.2),\n",
    "                                             nn.Linear(repr_dim * c_repr_base*2, c_repr_base*2),\n",
    "                                            )\n",
    "        elif c_repr_mode.startswith(\"c3\"):\n",
    "            self.class_embed = nn.Sequential(nn.Linear(repr_dim, repr_dim * c_repr_base*2),\n",
    "                                             nn.LeakyReLU(negative_slope=0.2),\n",
    "                                             nn.Linear(repr_dim * c_repr_base*2, repr_dim * c_repr_base*2),\n",
    "                                             nn.LeakyReLU(negative_slope=0.2),\n",
    "                                             nn.Linear(repr_dim * c_repr_base*2, c_repr_base*2),\n",
    "                                            )\n",
    "        elif c_repr_mode == \"None\":\n",
    "            pass\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        # Constructing self.z_embed:\n",
    "        self.z_mode = z_mode\n",
    "        self.z_dim = z_dim\n",
    "        if z_mode.startswith(\"c0\"):\n",
    "            def repeat_dim(z):\n",
    "                return torch.cat([z,z], 1)\n",
    "            self.z_embed = repeat_dim\n",
    "        elif z_mode.startswith(\"c1\"):\n",
    "            self.z_embed = nn.Linear(z_dim, z_dim*2)\n",
    "        elif z_mode.startswith(\"c2\"):\n",
    "            self.z_embed = nn.Sequential(nn.Linear(z_dim, z_dim*4),\n",
    "                                         nn.LeakyReLU(negative_slope=0.2),\n",
    "                                         nn.Linear(z_dim*4, z_dim*2),\n",
    "                                        )\n",
    "        elif z_mode.startswith(\"c3\"):\n",
    "            self.z_embed = nn.Sequential(nn.Linear(z_dim, z_dim*4),\n",
    "                                         nn.LeakyReLU(negative_slope=0.2),\n",
    "                                         nn.Linear(z_dim*4, z_dim*4),\n",
    "                                         nn.LeakyReLU(negative_slope=0.2),\n",
    "                                         nn.Linear(z_dim*4, z_dim*2),\n",
    "                                        )\n",
    "        elif z_mode == \"None\":\n",
    "            pass\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        self.skip = None\n",
    "\n",
    "        if in_channel != out_channel or downsample:\n",
    "            if is_spec_norm in [True, \"True\"]:\n",
    "                self.skip = nn.Sequential(\n",
    "                    spectral_norm(nn.Conv2d(in_channel, out_channel, 1, bias=False))\n",
    "                )\n",
    "            elif is_spec_norm in [False, \"False\"]:\n",
    "                self.skip = nn.Sequential(\n",
    "                    nn.Conv2d(in_channel, out_channel, 1, bias=False)\n",
    "                )\n",
    "            elif is_spec_norm == \"ws\":\n",
    "                self.skip = nn.Sequential(\n",
    "                    WSConv2d(in_channel, out_channel, 1, bias=False)\n",
    "                )\n",
    "\n",
    "        if self.downsample_mode != \"None\":\n",
    "            assert self.downsample_mode in [\"conv\", \"conv+rescale\"]\n",
    "            if self.downsample_mode == \"conv\":\n",
    "                self.conv_downsample = nn.Conv2d(out_channel, out_channel, kernel_size=kernel_size, padding=padding)\n",
    "            elif self.downsample_mode == \"conv+rescale\":\n",
    "                self.conv_downsample = nn.Conv2d(out_channel, out_channel*2, kernel_size=kernel_size, padding=padding)\n",
    "\n",
    "    def forward(self, input, c_repr=None, z=None):\n",
    "        out = input\n",
    "        if c_repr is not None:\n",
    "            if \"softmax\" in self.c_repr_mode:\n",
    "                c_repr = F.softmax(c_repr, dim=-1)\n",
    "            c_embed = self.class_embed(c_repr).view(input.shape[0], -1, 1, 1)  # [B, c_dim*4, 1, 1]\n",
    "        if z is not None:\n",
    "            assert (isinstance(z, tuple) or isinstance(z, list)) and len(z) == 1\n",
    "            z = z[0]\n",
    "            z_embed = self.z_embed(z).view(input.shape[0], -1, 1, 1)\n",
    "\n",
    "        if self.c_repr_mode == \"None\":\n",
    "            assert c_repr is None\n",
    "            if self.z_mode == \"None\":\n",
    "                out = self.conv1(out)\n",
    "                out = self.bn1(out)\n",
    "                out = self.act1(out)\n",
    "                if hasattr(self, \"dropout\") and self.dropout != 0:\n",
    "                    out = self.dropout_fn(out)\n",
    "                out = self.conv2(out)\n",
    "                out = self.bn2(out)\n",
    "            elif self.z_mode.startswith(\"c\"):\n",
    "                z_embed_1, z_embed_2 = z_embed.chunk(2, 1)\n",
    "                out_shape = out.shape[2:]\n",
    "                out = self.conv1(torch.cat([out, z_embed_1.expand(*z_embed_1.shape[:2], *out_shape)], 1))\n",
    "                out = self.bn1(out)\n",
    "                out = self.act1(out)\n",
    "                if hasattr(self, \"dropout\") and self.dropout != 0:\n",
    "                    out = self.dropout_fn(out)\n",
    "                out = self.conv2(torch.cat([out, z_embed_2.expand(*z_embed_2.shape[:2], *out_shape)], 1))\n",
    "                out = self.bn2(out)\n",
    "            else:\n",
    "                raise\n",
    "        elif self.c_repr_mode.startswith(\"l\"):\n",
    "            if self.z_mode == \"None\":\n",
    "                weight1, weight2, bias1, bias2 = c_embed.chunk(4, 1) # [B, out_channel, 1, 1]\n",
    "                out = self.conv1(out)    # [B, out_channel, H, W]\n",
    "                out = self.bn1(out)\n",
    "                out = weight1 * out + bias1\n",
    "                out = self.act1(out)\n",
    "                if hasattr(self, \"dropout\") and self.dropout != 0:\n",
    "                    out = self.dropout_fn(out)\n",
    "                out = self.conv2(out)\n",
    "                out = self.bn2(out)\n",
    "                out = weight2 * out + bias2\n",
    "            elif self.z_mode.startswith(\"c\"):\n",
    "                z_embed_1, z_embed_2 = z_embed.chunk(2, 1)\n",
    "                out_shape = out.shape[2:]\n",
    "                weight1, weight2, bias1, bias2 = c_embed.chunk(4, 1)\n",
    "                out = self.conv1(torch.cat([out, z_embed_1.expand(*z_embed_1.shape[:2], *out_shape)], 1))\n",
    "                out = self.bn1(out)\n",
    "                out = weight1 * out + bias1\n",
    "                out = self.act1(out)\n",
    "                if hasattr(self, \"dropout\") and self.dropout != 0:\n",
    "                    out = self.dropout_fn(out)\n",
    "                out = self.conv2(torch.cat([out, z_embed_2.expand(*z_embed_2.shape[:2], *out_shape)], 1))\n",
    "                out = self.bn2(out)\n",
    "                out = weight2 * out + bias2\n",
    "            else:\n",
    "                raise\n",
    "        elif self.c_repr_mode.startswith(\"c\"):\n",
    "            if self.z_mode == \"None\":\n",
    "                c_embed_1, c_embed_2 = c_embed.chunk(2, 1)\n",
    "                out_shape = out.shape[2:]\n",
    "                out = self.conv1(torch.cat([out, c_embed_1.expand(*c_embed_1.shape[:2], *out_shape)], 1))\n",
    "                out = self.bn1(out)\n",
    "                out = self.act1(out)\n",
    "                if hasattr(self, \"dropout\") and self.dropout != 0:\n",
    "                    out = self.dropout_fn(out)\n",
    "                out = self.conv2(torch.cat([out, c_embed_2.expand(*c_embed_2.shape[:2], *out_shape)], 1))\n",
    "                out = self.bn2(out)\n",
    "            elif self.z_mode.startswith(\"c\"):\n",
    "                c_embed_1, c_embed_2 = c_embed.chunk(2, 1)\n",
    "                z_embed_1, z_embed_2 = z_embed.chunk(2, 1)\n",
    "                out_shape = out.shape[2:]\n",
    "                out = self.conv1(torch.cat([out, z_embed_1.expand(*z_embed_1.shape[:2], *out_shape), c_embed_1.expand(*c_embed_1.shape[:2], *out_shape)], 1))\n",
    "                out = self.bn1(out)\n",
    "                out = self.act1(out)\n",
    "                if hasattr(self, \"dropout\") and self.dropout != 0:\n",
    "                    out = self.dropout_fn(out)\n",
    "                out = self.conv2(torch.cat([out, z_embed_2.expand(*z_embed_2.shape[:2], *out_shape), c_embed_2.expand(*c_embed_2.shape[:2], *out_shape)], 1))\n",
    "                out = self.bn2(out)\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        if hasattr(self, \"is_res\") and self.is_res:\n",
    "            if self.skip is not None:\n",
    "                skip = self.skip(input)\n",
    "            else:\n",
    "                skip = input\n",
    "            out = out + skip\n",
    "\n",
    "        if self.downsample:\n",
    "            if not hasattr(self, \"downsample_mode\") or self.downsample_mode == \"None\":\n",
    "                out = F.avg_pool2d(out, self.pool_kernel_size)\n",
    "            elif self.downsample_mode in [\"conv\", \"conv+rescale\"]:\n",
    "                out = self.conv_downsample(out)\n",
    "                out = F.avg_pool2d(out, self.pool_kernel_size)\n",
    "            else:\n",
    "                raise\n",
    "\n",
    "        out = self.act2(out)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2.4 Self-attention:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Self_Attn(nn.Module):\n",
    "    \"\"\" Self attention Layer\"\"\"\n",
    "    def __init__(self, in_dim, act_name):\n",
    "        super(Self_Attn, self).__init__()\n",
    "        self.chanel_in = in_dim\n",
    "        self.activation = get_activation(act_name)\n",
    "\n",
    "        self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim//8, kernel_size=1)\n",
    "        self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim//8, kernel_size=1)\n",
    "        self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)\n",
    "        self.gamma = nn.Parameter(torch.zeros(1))\n",
    "\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "            inputs :\n",
    "                x : input feature maps( B X C X H X W)\n",
    "            returns :\n",
    "                out : self attention value + input feature\n",
    "                attention: B x N x N (N is Width*Height)\n",
    "        \"\"\"\n",
    "        m_batchsize, C, width, height = x.size()\n",
    "        proj_query = self.query_conv(x).view(m_batchsize, -1, width*height).permute(0,2,1) # B x (N) x C\n",
    "        proj_key = self.key_conv(x).view(m_batchsize, -1, width*height) # B X C x (*H*W)\n",
    "        energy = torch.bmm(proj_query, proj_key) # B x N x N\n",
    "        attention = self.softmax(energy) # B x (N) x (N), for each of the N queries, perform softmax on the N keys\n",
    "        proj_value = self.value_conv(x).view(m_batchsize, -1, width*height) # B X C X N\n",
    "\n",
    "        out = torch.bmm(proj_value, attention.permute(0,2,1)) # value: BxCxN,  attention_perm: BxNxN: attention on the values.\n",
    "        out = out.view(m_batchsize, C, width, height)\n",
    "\n",
    "        out = self.gamma*out + x\n",
    "        return out, attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2.5 PositionEmbedding:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Adapted from https://github.com/vadimkantorov/yet_another_pytorch_slot_attention/blob/master/models.py\n",
    "class PositionEmbeddingImplicit(nn.Module):\n",
    "    def __init__(self, hidden_dim):\n",
    "        super().__init__()\n",
    "        self.dense = nn.Linear(4, hidden_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        spatial_shape = x.shape[-3:-1]\n",
    "        grid = torch.stack(torch.meshgrid(*[torch.linspace(0., 1., r, device = x.device) for r in spatial_shape]), dim = -1)\n",
    "        grid = torch.cat([grid, 1 - grid], dim = -1)\n",
    "        return x + self.dense(grid)\n",
    "\n",
    "class PositionEmbeddingSine(nn.Module):\n",
    "    # https://github.com/facebookresearch/detr/blob/master/models/position_encoding.py\n",
    "    def __init__(self, num_pos_feats=64, temperature=10000, normalize=True, scale=None):\n",
    "        super().__init__()\n",
    "        self.num_pos_feats = num_pos_feats\n",
    "        self.temperature = temperature\n",
    "        self.normalize = normalize\n",
    "        if scale is not None and normalize is False:\n",
    "            raise ValueError(\"normalize should be True if scale is passed\")\n",
    "        if scale is None:\n",
    "            scale = 2 * math.pi\n",
    "        self.scale = scale\n",
    "\n",
    "    def forward(self, x):\n",
    "        not_mask = torch.ones_like(x)\n",
    "        y_embed = not_mask.cumsum(1, dtype=torch.float32)\n",
    "        x_embed = not_mask.cumsum(2, dtype=torch.float32)\n",
    "        if self.normalize:\n",
    "            eps = 1e-6\n",
    "            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale\n",
    "            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale\n",
    "\n",
    "        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)\n",
    "        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)\n",
    "\n",
    "        pos_x = x_embed[:, :, :, None] / dim_t\n",
    "        pos_y = y_embed[:, :, :, None] / dim_t\n",
    "        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)\n",
    "        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)\n",
    "        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)\n",
    "        return x + pos\n",
    "\n",
    "\n",
    "class PositionEmbeddingLearned(nn.Module):\n",
    "    # https://github.com/facebookresearch/detr/blob/master/models/position_encoding.py\n",
    "    def __init__(self, num_pos_feats=256):\n",
    "        super().__init__()\n",
    "        self.row_embed = nn.Embedding(50, num_pos_feats)\n",
    "        self.col_embed = nn.Embedding(50, num_pos_feats)\n",
    "        #TODO: assert that x.shape matches the passed row_embed, col_embed\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        nn.init.uniform_(self.row_embed.weight)\n",
    "        nn.init.uniform_(self.col_embed.weight)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h, w = x.shape[-2:]\n",
    "        i = torch.arange(w, device=x.device)\n",
    "        j = torch.arange(h, device=x.device)\n",
    "        x_emb = self.col_embed(i)\n",
    "        y_emb = self.row_embed(j)\n",
    "        pos = torch.cat([\n",
    "            x_emb.unsqueeze(0).repeat(h, 1, 1),\n",
    "            y_emb.unsqueeze(1).repeat(1, w, 1),\n",
    "        ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)\n",
    "        return x + pos"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2.6 Downsample:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Downsample(nn.Module):\n",
    "    def __init__(self, pad_type='reflect', filt_size=3, stride=2, channels=None, pad_off=0):\n",
    "        super(Downsample, self).__init__()\n",
    "        self.filt_size = filt_size\n",
    "        self.pad_off = pad_off\n",
    "        self.pad_sizes = [int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)), int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2))]\n",
    "        self.pad_sizes = [pad_size+pad_off for pad_size in self.pad_sizes]\n",
    "        self.stride = stride\n",
    "        self.off = int((self.stride-1)/2.)\n",
    "        self.channels = channels\n",
    "\n",
    "        if(self.filt_size==1):\n",
    "            a = np.array([1.,])\n",
    "        elif(self.filt_size==2):\n",
    "            a = np.array([1., 1.])\n",
    "        elif(self.filt_size==3):\n",
    "            a = np.array([1., 2., 1.])\n",
    "        elif(self.filt_size==4):\n",
    "            a = np.array([1., 3., 3., 1.])\n",
    "        elif(self.filt_size==5):\n",
    "            a = np.array([1., 4., 6., 4., 1.])\n",
    "        elif(self.filt_size==6):\n",
    "            a = np.array([1., 5., 10., 10., 5., 1.])\n",
    "        elif(self.filt_size==7):\n",
    "            a = np.array([1., 6., 15., 20., 15., 6., 1.])\n",
    "\n",
    "        filt = torch.Tensor(a[:,None]*a[None,:])\n",
    "        filt = filt/torch.sum(filt)\n",
    "        self.register_buffer('filt', filt[None,None,:,:].repeat((self.channels,1,1,1)))\n",
    "\n",
    "        self.pad = get_pad_layer(pad_type)(self.pad_sizes)\n",
    "\n",
    "    def forward(self, inp):\n",
    "        if self.filt_size == 1:\n",
    "            if self.pad_off == 0:\n",
    "                return inp[:,:,::self.stride,::self.stride]\n",
    "            else:\n",
    "                return self.pad(inp)[:,:,::self.stride,::self.stride]\n",
    "        else:\n",
    "            return F.conv2d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1])\n",
    "\n",
    "\n",
    "def get_pad_layer(pad_type):\n",
    "    if(pad_type in ['refl','reflect']):\n",
    "        PadLayer = nn.ReflectionPad2d\n",
    "    elif(pad_type in ['repl','replicate']):\n",
    "        PadLayer = nn.ReplicationPad2d\n",
    "    elif(pad_type=='zero'):\n",
    "        PadLayer = nn.ZeroPad2d\n",
    "    else:\n",
    "        print('Pad type [%s] not recognized'%pad_type)\n",
    "    return PadLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     # BabyARC-relation dataset, 3D:\n",
    "#     from reasoning.experiments.concept_energy import get_dataset, ConceptDataset, ConceptDataset3D\n",
    "#     relation_args = init_args({\n",
    "#         \"dataset\": \"y-Parallel+VerticalMid+VerticalEdge\",\n",
    "#         \"seed\": 2,\n",
    "#         \"n_examples\": 40,\n",
    "#         \"canvas_size\": 16,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"1,2\",\n",
    "#         \"w_type\": \"image+mask\",\n",
    "#         \"color_map_3d\": \"same\",\n",
    "#         \"add_thick_surf\": (0, 0.5),\n",
    "#         \"add_thick_depth\": (0, 0.5),\n",
    "#         \"max_n_distractors\": 2,\n",
    "#         \"seed_3d\": 42,\n",
    "#         \"num_processes_3d\": 10,\n",
    "#         \"image_size_3d\": (256,256),\n",
    "#     })\n",
    "#     relation_dataset, args = get_dataset(relation_args, is_rewrite=False, is_load=True)\n",
    "#     relation_dataset.draw(range(40))\n",
    "#     tensor = relation_dataset[0][0][None]\n",
    "#     visualize_matrices([tensor[0]], use_color_dict=False)\n",
    "#     dd = Downsample(channels=3)\n",
    "#     tensor_d = dd(tensor)\n",
    "#     visualize_matrices([tensor_d[0]], use_color_dict=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.3 EBM Optimization:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3.1 SGLD ensemble:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def neg_mask_sgd_ensemble(\n",
    "    model,\n",
    "    img,\n",
    "    neg_mask,\n",
    "    c_repr,\n",
    "    z=None,\n",
    "    zgnn=None,\n",
    "    wtarget=None,\n",
    "    args=None,\n",
    "    ensemble_size=None,\n",
    "    out_mode=\"all\",\n",
    "    mask_info=None,\n",
    "    record_interval=-1,\n",
    "    return_history=False,\n",
    "    is_grad=False,\n",
    "    is_return_E=False,\n",
    "    batch_shape=None,\n",
    "):\n",
    "    \"\"\"\n",
    "    Perform an ensemble to discover a mask - run the model {ensemble_size}\n",
    "    times on the same image, and either return all the masks (out_mode) or\n",
    "    the one with the \"minimum\" energy.\n",
    "\n",
    "    The size of neg_mask and noise must be ensemble_size time that of\n",
    "    img and c_repr, but pos_img should not be duplicated manually.\n",
    "\n",
    "    Args:\n",
    "        record_interval: if set, indicates the intermediate mask should\n",
    "            be recorded every `record_interval` steps. Implies you should enable\n",
    "            `return_mask_history`.\n",
    "        return_history: whether the intermediate img/mask/z should be\n",
    "            returned at every step. If enabled, a third tuple element will be\n",
    "            returned.\n",
    "        is_grad: if True, will use neg_mask_sgd_with_kl. Otherwise use neg_mask_sgd.\n",
    "        batch_shape: if not None, will be (B_task, B_example).\n",
    "        mask_info: dictionary for various constraint for mask. The key includes:\n",
    "            \"mask_exclude\": if its value is not None, will have energy for the overlapping between mask_exclude and neg_mask, \n",
    "                with coefficient of SGLD_mutual_exclusive_coef.\n",
    "    \"\"\"\n",
    "    # repeat ensemble_size time along the batch dimension:\n",
    "    img_l, c_repr_l = repeat_n(img, c_repr, n_repeats=ensemble_size)\n",
    "    if mask_info is not None and \"mask_exclude\" in mask_info:\n",
    "        mask_info = deepcopy(mask_info)\n",
    "        mask_info[\"mask_exclude\"] = repeat_n(mask_info[\"mask_exclude\"], n_repeats=ensemble_size)\n",
    "\n",
    "    # Find neg_mask for the batch:\n",
    "    if is_grad:\n",
    "        _, (img, neg_mask, c_repr, z, zgnn, wtarget), info = neg_mask_sgd_with_kl(\n",
    "            model,\n",
    "            img=img_l, neg_mask=neg_mask, c_repr=c_repr_l, z=z, zgnn=zgnn, wtarget=wtarget,\n",
    "            args=args,\n",
    "            mask_info=mask_info,\n",
    "            is_return_E=is_return_E,\n",
    "            batch_shape=batch_shape,\n",
    "            record_interval=record_interval,\n",
    "        )\n",
    "    else:\n",
    "        (img, neg_mask, c_repr, z, zgnn, wtarget), info = neg_mask_sgd(\n",
    "            model,\n",
    "            img=img_l, neg_mask=neg_mask, c_repr=c_repr_l, z=z, zgnn=zgnn, wtarget=wtarget,\n",
    "            args=args,\n",
    "            mask_info=mask_info,\n",
    "            is_return_E=is_return_E,\n",
    "            batch_shape=batch_shape,\n",
    "            record_interval=record_interval,\n",
    "        )\n",
    "\n",
    "    for key in info:\n",
    "        if key.endswith(\"_list\") and key not in [\"img_list\", \"neg_mask_list\", \"c_repr_list\", \"z_list\", \"zgnn_list\", \"wtarget_list\"]:\n",
    "            info[key] = info[key].reshape(len(info[key]), ensemble_size, -1)  # [sample_step, ensemble_size, batch_size]\n",
    "        else:\n",
    "            if isinstance(info[key], tuple) or isinstance(info[key], list):\n",
    "                if len(info[key]) > 0:\n",
    "                    shape_rest = info[key][0].shape[2:]\n",
    "                    info[key] = tuple(info[key][k].reshape(len(info[key][k]), ensemble_size, -1, *shape_rest) for k in range(len(info[key])))\n",
    "            else:\n",
    "                shape_rest = info[key].shape[2:]\n",
    "                info[key] = info[key].reshape(len(info[key]), ensemble_size, -1, *shape_rest)\n",
    "\n",
    "    neg_out_list_reshape = info.pop(\"neg_out_list\")\n",
    "\n",
    "    if \"image\" in args.ebm_target.split(\"+\"):\n",
    "        if args.is_image_tuple:\n",
    "            img_reshape = tuple(img[k].reshape(\n",
    "                ensemble_size, -1, *img[k].shape[1:]) for k in range(len(img)))  # [ensemble_size, batch_size, C, H, W]\n",
    "        else:\n",
    "            img_reshape = img.reshape(ensemble_size, -1, *img.shape[1:])  # [ensemble_size, batch_size, C, H, W]\n",
    "    else:\n",
    "        img_reshape = None\n",
    "    if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "        neg_mask_reshape = tuple(neg_mask[k].reshape(\n",
    "            ensemble_size, -1, *neg_mask[k].shape[1:]) for k in range(len(neg_mask)))  # [ensemble_size, batch_size, C, H, W]\n",
    "    else:\n",
    "        neg_mask_reshape = None\n",
    "    if \"z\" in args.ebm_target.split(\"+\"):\n",
    "        z_reshape = tuple(z[k].reshape(\n",
    "            ensemble_size, -1, *z[k].shape[1:]) if z[k] is not None else None for k in range(len(z)))  # Each [ensemble_size, batch_size, Z]\n",
    "    else:\n",
    "        z_reshape = None\n",
    "    if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "        assert ensemble_size == 1\n",
    "        zgnn_reshape = tuple(zgnn[k].reshape(\n",
    "            ensemble_size, -1, *zgnn[k].shape[1:]) if zgnn[k] is not None else None for k in range(len(zgnn)))  # Each [ensemble_size, batch_size, Z]\n",
    "    else:\n",
    "        zgnn_reshape = None\n",
    "    if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "        assert ensemble_size == 1\n",
    "        wtarget_reshape = wtarget.reshape(ensemble_size, -1, *wtarget.shape[1:])  # [ensemble_size, batch_size, C, H, W]\n",
    "    else:\n",
    "        wtarget_reshape = None\n",
    "\n",
    "    if out_mode == \"min\":\n",
    "        idx_argmin = neg_out_list_reshape[-1].argmin(0)  # [B,] each number has value up to ensemble_size-1\n",
    "        if \"image\" in args.ebm_target.split(\"+\"):\n",
    "            if args.is_image_tuple:\n",
    "                img_reshape = tuple(gather_broadcast(img_reshape[k].transpose(\n",
    "                    0, 1), 1, idx_argmin) for k in range(len(img)))\n",
    "            else:\n",
    "                img_reshape = gather_broadcast(img_reshape.transpose(\n",
    "                    0, 1), 1, idx_argmin)\n",
    "        if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            neg_mask_reshape = tuple(gather_broadcast(neg_mask_reshape[k].transpose(\n",
    "                0, 1), 1, idx_argmin) for k in range(len(neg_mask)))\n",
    "        if \"z\" in args.ebm_target.split(\"+\"):\n",
    "            z_reshape = tuple(gather_broadcast(z_reshape[k].transpose(\n",
    "                0, 1), 1, idx_argmin) if z[k] is not None else None for k in range(len(z)))\n",
    "        if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "            zgnn_reshape = tuple(gather_broadcast(zgnn_reshape[k].transpose(\n",
    "                0, 1), 1, idx_argmin) if zgnn[k] is not None else None for k in range(len(zgnn)))\n",
    "        if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "            wtarget_reshape = gather_broadcast(wtarget_reshape.transpose(\n",
    "                0, 1), 1, idx_argmin)\n",
    "    elif out_mode == \"all\":\n",
    "        pass\n",
    "    elif out_mode == \"all-sorted\":\n",
    "        idx_argsort = torch.LongTensor(neg_out_list_reshape[-1].argsort(0))  # [ensemble_size, B] each number has value up to ensemble_size-1\n",
    "        # neg_out_list_reshape (shape of [sample_step, ensemble_size, B]):\n",
    "        index = torch.LongTensor(idx_argsort[None]).expand(neg_out_list_reshape.shape).numpy()\n",
    "        neg_out_list_reshape = np.take_along_axis(neg_out_list_reshape, indices=index, axis=1)\n",
    "        # Other results:\n",
    "        if \"image\" in args.ebm_target.split(\"+\"):\n",
    "            if args.is_image_tuple:\n",
    "                index = idx_argsort[:,:,None,None,None].to(img_reshape[0].device).expand_as(img_reshape[0])  # [ensemble_size, B, ...]\n",
    "                img_reshape = tuple(torch.gather(img_reshape[k], dim=0, index=index) for k in range(len(img)))\n",
    "            else:\n",
    "                index = idx_argsort[:,:,None,None,None].to(img_reshape.device).expand_as(img_reshape)\n",
    "                img_reshape = torch.gather(img_reshape, dim=0, index=index)\n",
    "        if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            index = idx_argsort[:,:,None,None,None].to(neg_mask_reshape[0].device).expand_as(neg_mask_reshape[0])  # [ensemble_size, B, ...]\n",
    "            neg_mask_reshape = tuple(torch.gather(neg_mask_reshape[k], dim=0, index=index) for k in range(len(neg_mask_reshape)))\n",
    "        if \"z\" in args.ebm_target.split(\"+\"):\n",
    "            index = idx_argsort[:,:,None].to(z_reshape[0].device).expand_as(z_reshape[0])\n",
    "            z_reshape = tuple(torch.gather(z_reshape[k], dim=0, index=index) if z[k] is not None else None for k in range(len(z)))\n",
    "        if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "            index = idx_argsort[:,:,None].to(zgnn_reshape[0].device).expand_as(zgnn_reshape[0])\n",
    "            zgnn_reshape = tuple(torch.gather(zgnn_reshape[k], dim=0, index=index) if zgnn[k] is not None else None for k in range(len(zgnn)))\n",
    "        if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "            index = idx_argsort[:,:,None,None,None].to(wtarget_reshape.device).expand_as(wtarget_reshape)\n",
    "            wtarget_reshape = torch.gather(wtarget_reshape, dim=0, index=index)\n",
    "        # model.info:\n",
    "        if hasattr(model, \"info\"):\n",
    "            for key, value in model.info.items():\n",
    "                index = idx_argsort.numpy()  # [ensemble_size, B]\n",
    "                model.info[key] = np.take_along_axis(value, indices=index, axis=0)\n",
    "        # info:\n",
    "        for key, value in info.items():\n",
    "            if key.endswith(\"_list\"):\n",
    "                if key not in [\"img_list\", \"neg_mask_list\", \"c_repr_list\", \"z_list\", \"zgnn_list\", \"wtarget_list\"]:\n",
    "                    # The value has shape [sample_step, ensemble_size, B]:\n",
    "                    index = torch.LongTensor(idx_argsort[None]).expand(value.shape).numpy()\n",
    "                    info[key] = np.take_along_axis(value, indices=index, axis=1)\n",
    "                elif len(value) > 0:\n",
    "                    if not isinstance(value, tuple) and not isinstance(value, list):\n",
    "                        index = torch.LongTensor(extend_dims(idx_argsort[None], n_dims=len(value.shape), loc=\"right\")).expand(value.shape).numpy()\n",
    "                        info[key] = np.take_along_axis(value, indices=index, axis=1)\n",
    "                    else:\n",
    "                        index = torch.LongTensor(extend_dims(idx_argsort[None], n_dims=len(value[0].shape), loc=\"right\")).expand(value[0].shape).numpy()\n",
    "                        info[key] = tuple(np.take_along_axis(value_ele, indices=index, axis=1) for value_ele in value)\n",
    "    else:\n",
    "        raise\n",
    "\n",
    "    if return_history or is_return_E:\n",
    "        return (img_reshape, neg_mask_reshape, z_reshape, zgnn_reshape, wtarget_reshape), neg_out_list_reshape, info\n",
    "    else:\n",
    "        return (img_reshape, neg_mask_reshape, z_reshape, zgnn_reshape, wtarget_reshape), neg_out_list_reshape, {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3.2 SGLD:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def neg_mask_sgd(\n",
    "    model,\n",
    "    img=None,\n",
    "    neg_mask=None,\n",
    "    c_repr=None,\n",
    "    z=None,\n",
    "    zgnn=None,\n",
    "    wtarget=None,\n",
    "    args=None,\n",
    "    mask_info=None,\n",
    "    is_return_E=False,\n",
    "    batch_shape=None,\n",
    "    record_interval=-1,\n",
    "):\n",
    "    \"\"\"Perform SGLD w.r.t. a subset of {img, neg_mask, c_repr, z} given the others.\"\"\"\n",
    "    requires_grad(model.parameters(), False)\n",
    "    model.eval()\n",
    "    neg_out_list = []\n",
    "    if args.step_size_img == -1:\n",
    "        args.step_size_img = args.step_size\n",
    "    if args.step_size_z == -1:\n",
    "        args.step_size_z = args.step_size\n",
    "    if args.step_size_zgnn == -1:\n",
    "        args.step_size_zgnn = args.step_size\n",
    "    if args.step_size_wtarget == -1:\n",
    "        args.step_size_wtarget = args.step_size\n",
    "\n",
    "    # Initialize variables (a subset of {img, neg_mask, c_repr, z}):\n",
    "    if img is not None:\n",
    "        batch_size = img[0].shape[0] if args.is_image_tuple else img.shape[0]\n",
    "        in_channels = img[0].shape[1] if args.is_image_tuple else img.shape[1]\n",
    "        device = img[0].device if args.is_image_tuple else img.device\n",
    "        img_shape = img[0].shape[2:] if args.is_image_tuple else img.shape[2:]\n",
    "    else:\n",
    "        batch_size = neg_mask[0].shape[0]\n",
    "        in_channels = args.in_channels\n",
    "        img_shape = args.image_size if args.rescaled_size == \"None\" else eval(args.rescaled_size)\n",
    "        device = neg_mask[0].device\n",
    "\n",
    "    img_list = []\n",
    "    if \"image\" in args.ebm_target.split(\"+\"):\n",
    "        img_value_min, img_value_max = args.image_value_range.split(\",\")\n",
    "        img_value_min, img_value_max = eval(img_value_min), eval(img_value_max)\n",
    "        img_value_span = img_value_max - img_value_min\n",
    "        assert img_value_span >= 1\n",
    "        if img is None:\n",
    "            img = (\n",
    "                torch.rand(batch_size, in_channels, *img_shape, device=device) * img_value_span + img_value_min, torch.rand(batch_size, in_channels, *img_shape, device=device) * img_value_span + img_value_min\n",
    "            ) if args.is_image_tuple else torch.rand(batch_size, in_channels, *img_shape, device=device) * img_value_span + img_value_min\n",
    "        else:\n",
    "            if isinstance(img, tuple):\n",
    "                img = tuple(torch.rand(batch_size, in_channels, *img_shape, device=device) * img_value_span + img_value_min if img_ele is None else img_ele for img_ele in img)\n",
    "        if args.is_image_tuple:\n",
    "            for i in range(len(img)):\n",
    "                img[i].requires_grad = True\n",
    "        else:\n",
    "            img.requires_grad = True\n",
    "        if record_interval != -1:\n",
    "            img_list.append(deepcopy(tuple(to_np_array(*img, keep_list=True)) if args.is_image_tuple else to_np_array(img)))\n",
    "\n",
    "    neg_mask_list = []\n",
    "    if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "        if neg_mask is None:\n",
    "            img_shape = img[0].shape[2:] if args.is_image_tuple else img.shape[2:]\n",
    "            w_dim = 1 if \"mask\" in model.w_type else in_channels\n",
    "            neg_mask = tuple(torch.rand(batch_size, w_dim, *img_shape, device=device) for _ in range(model.mask_arity))\n",
    "        for i in range(model.mask_arity):\n",
    "            neg_mask[i].requires_grad = True\n",
    "        if record_interval != -1:\n",
    "            neg_mask_list.append(deepcopy(tuple(to_np_array(*neg_mask, keep_list=True))))\n",
    "\n",
    "    c_repr_list = []\n",
    "    if \"repr\" in args.ebm_target.split(\"+\"):\n",
    "        c_repr = torch.rand(batch_size, REPR_DIM, device=device) if c_repr is None else c_repr.to(device)\n",
    "        c_repr.requires_grad = True\n",
    "        if record_interval != -1:\n",
    "            c_repr_list.append(deepcopy(to_np_array(c_repr)))\n",
    "\n",
    "    z_list = []\n",
    "    if \"z\" in args.ebm_target.split(\"+\"):\n",
    "        z_len = 1\n",
    "        z = tuple(torch.rand(batch_size, model.z_dim, device=device) for _ in range(z_len)) if z is None else tuple(to_device_recur(z_ele, device) for z_ele in z)\n",
    "        z_len = len(z)\n",
    "        for i in range(z_len):\n",
    "            if z[i] is not None:\n",
    "                z[i].requires_grad = True\n",
    "        if record_interval != -1:\n",
    "            z_list.append(deepcopy(tuple(to_np_array(*z, keep_list=True))))\n",
    "    assert z is None or isinstance(z, tuple) or isinstance(z, list)\n",
    "\n",
    "    zgnn_list = []\n",
    "    if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "        n_nodes = len(z)\n",
    "        n_edges = n_nodes * (n_nodes - 1)\n",
    "        zgnn_dim = model.zgnn_dim\n",
    "        zgnn = (torch.rand(batch_shape[0], n_nodes, zgnn_dim, device=device) if model.gnn.is_zgnn_node else None, torch.rand(batch_shape[0], n_edges, model.edge_attr_size, device=device)) if zgnn is None else tuple(to_device_recur(zgnn_ele, device) for zgnn_ele in zgnn)\n",
    "        zgnn_len = len(zgnn)\n",
    "        for i in range(zgnn_len):\n",
    "            if zgnn[i] is not None:\n",
    "                zgnn[i].requires_grad = True\n",
    "        if record_interval != -1:\n",
    "            zgnn_list.append(deepcopy(tuple(to_np_array(*zgnn, keep_list=True))))\n",
    "\n",
    "    wtarget_list = []\n",
    "    if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "        w_dim = 1 if \"mask\" in model.w_type else in_channels\n",
    "        if wtarget is None:\n",
    "            wtarget = torch.rand(batch_size, w_dim, *img_shape, device=device)\n",
    "        wtarget.requires_grad = True\n",
    "        if record_interval != -1:\n",
    "            wtarget_list.append(deepcopy(to_np_array(wtarget)))\n",
    "\n",
    "    # Setting up noise and step_size scheduling:\n",
    "    if args.lambd_start == -1:\n",
    "        args.lambd_start = args.lambd\n",
    "    lambd_list = args.lambd + 1/2 * (args.lambd_start - args.lambd) * (1 + torch.cos(torch.arange(args.sample_step)/args.sample_step * np.pi))\n",
    "    if args.step_size_start == -1:\n",
    "        args.step_size_start = args.step_size\n",
    "    step_size_list = args.step_size + 1/2 * (args.step_size_start - args.step_size) * (1 + torch.cos(torch.arange(args.sample_step)/args.sample_step * np.pi))\n",
    "    if args.SGLD_is_anneal:\n",
    "        multiplier_list = np.linspace(0, 1, args.sample_step) ** args.SGLD_anneal_power\n",
    "    else:\n",
    "        multiplier_list = np.ones(args.sample_step)\n",
    "\n",
    "    if args.SGLD_object_exceed_coef > 0:\n",
    "        if isinstance(img, tuple):\n",
    "            pos_all_mask = ((img[0][:,:1] == 1).float(), (img[1][:,:1] == 1).float())\n",
    "        else:\n",
    "            pos_all_mask = (img[:,:1] == 1).float()\n",
    "\n",
    "    reg_dict = {}\n",
    "\n",
    "    # SGLD:\n",
    "    for k in range(args.sample_step):\n",
    "        multiplier = multiplier_list[k]\n",
    "        # Each step add noise using Langevin dynamics:\n",
    "        if \"image\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise_img\" not in locals():\n",
    "                noise_img = (torch.randn(batch_size, in_channels, *img_shape, device=device), torch.randn(batch_size, in_channels, *img_shape, device=device)) if args.is_image_tuple else torch.randn(batch_size, in_channels, *img_shape, device=device)\n",
    "            if lambd_list[k] > 0:\n",
    "                if args.is_image_tuple:\n",
    "                    noise_img[0].normal_(0, lambd_list[k])\n",
    "                    noise_img[1].normal_(0, lambd_list[k])\n",
    "                    img[0].data.add_(noise_img[0].data)\n",
    "                    img[1].data.add_(noise_img[1].data)\n",
    "                else:\n",
    "                    noise_img.normal_(0, lambd_list[k])\n",
    "                    img.data.add_(noise_img.data)\n",
    "        if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise\" not in locals():\n",
    "                w_dim = 1 if \"mask\" in model.w_type else in_channels\n",
    "                noise = tuple(torch.randn(batch_size, w_dim, *(args.image_size if args.rescaled_size == \"None\" else eval(args.rescaled_size)), device=device) for _ in range(model.mask_arity))\n",
    "            if lambd_list[k] > 0:\n",
    "                for i in range(model.mask_arity):\n",
    "                    noise[i].normal_(0, lambd_list[k])\n",
    "                    neg_mask[i].data.add_(noise[i].data)\n",
    "        if \"repr\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise_repr\" not in locals():\n",
    "                noise_repr = torch.randn(batch_size, REPR_DIM, device=device)\n",
    "            if lambd_list[k] > 0:\n",
    "                noise_repr.normal_(0, lambd_list[k])\n",
    "                c_repr.data.add_(noise_repr.data)\n",
    "        if \"z\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise_z\" not in locals():\n",
    "                noise_z = tuple(torch.randn(batch_size, model.z_dim, device=device) if z[i] is not None else None for i in range(z_len))\n",
    "            if lambd_list[k] > 0:\n",
    "                for i in range(z_len):\n",
    "                    if z[i] is not None:\n",
    "                        noise_z[i].normal_(0, lambd_list[k])\n",
    "                        z[i].data.add_(noise_z[i].data)\n",
    "        if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise_zgnn\" not in locals():\n",
    "                noise_zgnn = tuple(torch.randn(zgnn[i].shape, device=device) if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "            if lambd_list[k] > 0:\n",
    "                for i in range(zgnn_len):\n",
    "                    if zgnn[i] is not None:\n",
    "                        noise_zgnn[i].normal_(0, lambd_list[k])\n",
    "                        zgnn[i].data.add_(noise_zgnn[i].data)\n",
    "        if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise_wtarget\" not in locals():\n",
    "                noise_wtarget = torch.randn(wtarget.shape, device=device)\n",
    "            if lambd_list[k] > 0:\n",
    "                noise_wtarget.normal_(0, lambd_list[k])\n",
    "                wtarget.data.add_(noise_wtarget.data)\n",
    "\n",
    "        # Compute the energy for each (pos_image, neg_mask, concept) instance:\n",
    "        if k == args.sample_step - 1 and is_return_E:\n",
    "            neg_out = model(img, neg_mask, c_repr=c_repr, z=z, zgnn=zgnn, wtarget=wtarget, batch_shape=batch_shape, is_E_tensor=is_return_E)\n",
    "            E_all = model.info[\"E_all\"]\n",
    "        else:\n",
    "            neg_out = model(img, neg_mask, c_repr=c_repr, z=z, zgnn=zgnn, wtarget=wtarget, batch_shape=batch_shape)\n",
    "        neg_out_sum = neg_out.sum()\n",
    "        record_data(reg_dict, deepcopy(to_np_array(neg_out)), \"neg_out_core_list\")\n",
    "        if hasattr(model, \"info\"):\n",
    "            for key in model.info:\n",
    "                record_data(reg_dict, deepcopy(to_np_array(model.info[key])), \"{}_list\".format(key))\n",
    "\n",
    "        # Penalize areas where the masks overlap with each other\n",
    "        if args.SGLD_mutual_exclusive_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            mutual_exclusive = get_neg_mask_overlap(neg_mask, mask_info=mask_info, is_penalize_lower=args.SGLD_is_penalize_lower, img=img) * args.SGLD_mutual_exclusive_coef * multiplier\n",
    "            neg_out_sum += mutual_exclusive.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(mutual_exclusive)), \"mutual_exclusive_list\")\n",
    "        if args.SGLD_fine_mutual_exclusive_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            fine_mutual_exclusive = get_fine_neg_mask_overlap(neg_mask, mask_info=mask_info) * args.SGLD_fine_mutual_exclusive_coef * multiplier\n",
    "            neg_out_sum += fine_mutual_exclusive.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(fine_mutual_exclusive)), \"fine_mutual_exclusive_list\")\n",
    "        # Penalize areas where the negative mask hits areas not covered by the true mask\n",
    "        if args.SGLD_object_exceed_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            # SGLD_object_exceed_coef is annealed quadratically as the sample steps continue\n",
    "            object_exceed = get_neg_mask_exceed(neg_mask, pos_all_mask) * args.SGLD_object_exceed_coef * multiplier\n",
    "            neg_out_sum += object_exceed.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(object_exceed)), \"object_exceed_list\")\n",
    "        # Encourage each EBM will either fully explain or fully ignore each pixel:\n",
    "        if args.SGLD_pixel_entropy_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            pixel_entropy = get_pixel_entropy(neg_mask) * args.SGLD_pixel_entropy_coef * multiplier\n",
    "            neg_out_sum += pixel_entropy.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(pixel_entropy)), \"pixel_entropy_list\")\n",
    "        # Encourage each EBM mask will either fully explain an object or not explain any\n",
    "        if args.SGLD_mask_entropy_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            mask_entropy = get_mask_entropy(neg_mask) * args.SGLD_mask_entropy_coef * multiplier\n",
    "            neg_out_sum += mask_entropy.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(mask_entropy)), \"mask_entropy_list\")\n",
    "        # Encourage the EBMs to specialize, where the ebms whose masks that are nearest to 0 or 1 will be push hardes towards 0 or 1:\n",
    "        if args.SGLD_pixel_gm_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            pixel_gm = get_pixel_gm(neg_mask) * args.SGLD_pixel_gm_coef * multiplier\n",
    "            neg_out_sum += pixel_gm.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(pixel_gm)), \"pixel_gm_list\")\n",
    "        if \"mask\" in args.ebm_target.split(\"+\") and (\n",
    "            args.SGLD_iou_batch_consistency_coef > 0 or\n",
    "            args.SGLD_iou_concept_repel_coef > 0 or\n",
    "            args.SGLD_iou_relation_repel_coef > 0 or\n",
    "            args.SGLD_iou_relation_overlap_coef > 0 or\n",
    "            args.SGLD_iou_attract_coef > 0):\n",
    "            graph_energy = get_graph_energy(\n",
    "                neg_mask,\n",
    "                mask_info=mask_info,\n",
    "                iou_batch_consistency_coef=args.SGLD_iou_batch_consistency_coef,\n",
    "                iou_concept_repel_coef=args.SGLD_iou_concept_repel_coef,\n",
    "                iou_relation_repel_coef=args.SGLD_iou_relation_repel_coef,\n",
    "                iou_relation_overlap_coef=args.SGLD_iou_relation_overlap_coef,\n",
    "                iou_attract_coef=args.SGLD_iou_attract_coef,\n",
    "                batch_shape=batch_shape,\n",
    "            )[0] * multiplier\n",
    "            neg_out_sum += graph_energy.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(graph_energy)), \"graph_energy_list\")\n",
    "\n",
    "        neg_out_sum.backward()\n",
    "\n",
    "        # Perform gradient descent on the neg_mask(s) and/or c_repr:\n",
    "        if \"image\" in args.ebm_target.split(\"+\"):\n",
    "            if args.is_image_tuple:\n",
    "                for i in range(len(img)):\n",
    "                    img[i].grad.data.clamp_(-0.01, 0.01)\n",
    "                    img[i].data.add_(img[i].grad.data, alpha=-args.step_size_img)\n",
    "                    img[i].grad.detach_()\n",
    "                    img[i].grad.zero_()\n",
    "                    img[i].data.clamp_(img_value_min, img_value_max)\n",
    "                if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                    img_list.append(deepcopy(tuple(to_np_array(*img, keep_list=True))))\n",
    "            else:\n",
    "                img.grad.data.clamp_(-0.01, 0.01)\n",
    "                img.data.add_(img.grad.data, alpha=-args.step_size_img)\n",
    "                img.grad.detach_()\n",
    "                img.grad.zero_()\n",
    "                img.data.clamp_(img_value_min, img_value_max)\n",
    "                if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                    img_list.append(deepcopy(to_np_array(img)))\n",
    "        if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            for i in range(model.mask_arity):\n",
    "                neg_mask[i].grad.data.clamp_(-0.01, 0.01)\n",
    "                neg_mask[i].data.add_(neg_mask[i].grad.data, alpha=-step_size_list[k])\n",
    "                neg_mask[i].grad.detach_()\n",
    "                neg_mask[i].grad.zero_()\n",
    "                neg_mask[i].data.clamp_(0, 1)\n",
    "            if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                neg_mask_list.append(deepcopy(tuple(to_np_array(*neg_mask, keep_list=True))))\n",
    "        if \"repr\" in args.ebm_target.split(\"+\"):\n",
    "            c_repr.grad.data.clamp_(-0.01, 0.01)\n",
    "            c_repr.data.add_(c_repr.grad.data, alpha=-args.step_size_repr)\n",
    "            c_repr.grad.detach_()\n",
    "            c_repr.grad.zero_()\n",
    "            if \"softmax\" not in args.c_repr_mode:\n",
    "                c_repr.data.clamp_(0, 1)\n",
    "            if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                c_repr_list.append(deepcopy(to_np_array(c_repr)))\n",
    "        if \"z\" in args.ebm_target.split(\"+\"):\n",
    "            for i in range(len(z)):\n",
    "                if z[i] is not None:\n",
    "                    z[i].grad.data.clamp_(-0.01, 0.01)\n",
    "                    z[i].data.add_(z[i].grad.data, alpha=-args.step_size_z)\n",
    "                    z[i].grad.detach_()\n",
    "                    z[i].grad.zero_()\n",
    "                    z[i].data.clamp_(0, 1)\n",
    "            if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                z_list.append(deepcopy(tuple(to_np_array(*z, keep_list=True))))\n",
    "        if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "            for i in range(zgnn_len):\n",
    "                if zgnn[i] is not None:\n",
    "                    zgnn[i].grad.data.clamp_(-0.01, 0.01)\n",
    "                    zgnn[i].data.add_(zgnn[i].grad.data, alpha=-args.step_size_zgnn)\n",
    "                    zgnn[i].grad.detach_()\n",
    "                    zgnn[i].grad.zero_()\n",
    "                    zgnn[i].data.clamp_(0, 1)\n",
    "            if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                zgnn_list.append(deepcopy(tuple(to_np_array(*zgnn, keep_list=True))))\n",
    "        if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "            wtarget.grad.data.clamp_(-0.01, 0.01)\n",
    "            wtarget.data.add_(wtarget.grad.data, alpha=-args.step_size_wtarget)\n",
    "            wtarget.grad.detach_()\n",
    "            wtarget.grad.zero_()\n",
    "            wtarget.data.clamp_(0, 1)\n",
    "            if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                wtarget_list.append(deepcopy(tuple(to_np_array(*wtarget, keep_list=True))))\n",
    "\n",
    "        neg_out_list.append(deepcopy(to_np_array(neg_out)))\n",
    "\n",
    "    \"\"\"\n",
    "    Record the energies of the last step:\n",
    "    \"\"\"\n",
    "    neg_out = model(img, neg_mask, c_repr=c_repr, z=z, zgnn=zgnn, wtarget=wtarget, batch_shape=batch_shape)\n",
    "    neg_out_list.append(deepcopy(to_np_array(neg_out)))\n",
    "    if hasattr(model, \"info\"):\n",
    "        for key in model.info:\n",
    "            record_data(reg_dict, deepcopy(to_np_array(model.info[key])), \"{}_list\".format(key))\n",
    "    # Penalize areas where the masks overlap with each other\n",
    "    if args.SGLD_mutual_exclusive_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "        mutual_exclusive = get_neg_mask_overlap(neg_mask, mask_info=mask_info, is_penalize_lower=args.SGLD_is_penalize_lower, img=img) * args.SGLD_mutual_exclusive_coef * multiplier\n",
    "        record_data(reg_dict, deepcopy(to_np_array(mutual_exclusive)), \"mutual_exclusive_list\")\n",
    "    if args.SGLD_fine_mutual_exclusive_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "        fine_mutual_exclusive = get_fine_neg_mask_overlap(neg_mask, mask_info=mask_info) * args.SGLD_fine_mutual_exclusive_coef * multiplier\n",
    "        record_data(reg_dict, deepcopy(to_np_array(fine_mutual_exclusive)), \"fine_mutual_exclusive_list\")\n",
    "    # Penalize areas where the negative mask hits areas not covered by the true mask\n",
    "    if args.SGLD_object_exceed_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "        # SGLD_object_exceed_coef is annealed quadratically as the sample steps continue\n",
    "        object_exceed = get_neg_mask_exceed(neg_mask, pos_all_mask) * args.SGLD_object_exceed_coef * multiplier\n",
    "        record_data(reg_dict, deepcopy(to_np_array(object_exceed)), \"object_exceed_list\")\n",
    "    # Encourage each EBM will either fully explain or fully ignore each pixel:\n",
    "    if args.SGLD_pixel_entropy_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "        pixel_entropy = get_pixel_entropy(neg_mask) * args.SGLD_pixel_entropy_coef * multiplier\n",
    "        record_data(reg_dict, deepcopy(to_np_array(pixel_entropy)), \"pixel_entropy_list\")\n",
    "    # Encourage each EBM mask will either fully explain an object or not explain any\n",
    "    if args.SGLD_mask_entropy_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "        mask_entropy = get_mask_entropy(neg_mask) * args.SGLD_mask_entropy_coef * multiplier\n",
    "        record_data(reg_dict, deepcopy(to_np_array(mask_entropy)), \"mask_entropy_list\")\n",
    "    # Encourage the EBMs to specialize, where the ebms whose masks that are nearest to 0 or 1 will be push hardes towards 0 or 1:\n",
    "    if args.SGLD_pixel_gm_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "        pixel_gm = get_pixel_gm(neg_mask) * args.SGLD_pixel_gm_coef * multiplier\n",
    "        record_data(reg_dict, deepcopy(to_np_array(pixel_gm)), \"pixel_gm_list\")\n",
    "    if \"mask\" in args.ebm_target.split(\"+\") and (\n",
    "        args.SGLD_iou_batch_consistency_coef > 0 or\n",
    "        args.SGLD_iou_concept_repel_coef > 0 or\n",
    "        args.SGLD_iou_relation_repel_coef > 0 or\n",
    "        args.SGLD_iou_relation_overlap_coef > 0 or\n",
    "        args.SGLD_iou_attract_coef > 0):\n",
    "        graph_energy = get_graph_energy(\n",
    "            neg_mask,\n",
    "            mask_info=mask_info,\n",
    "            iou_batch_consistency_coef=args.SGLD_iou_batch_consistency_coef,\n",
    "            iou_concept_repel_coef=args.SGLD_iou_concept_repel_coef,\n",
    "            iou_relation_repel_coef=args.SGLD_iou_relation_repel_coef,\n",
    "            iou_relation_overlap_coef=args.SGLD_iou_relation_overlap_coef,\n",
    "            iou_attract_coef=args.SGLD_iou_attract_coef,\n",
    "            batch_shape=batch_shape,\n",
    "        )[0] * multiplier\n",
    "        record_data(reg_dict, deepcopy(to_np_array(graph_energy)), \"graph_energy_list\")\n",
    "\n",
    "    neg_out_list = np.concatenate(neg_out_list, -1).T   # after: [sample_step, B]\n",
    "\n",
    "    if \"image\" in args.ebm_target.split(\"+\"):\n",
    "        if args.is_image_tuple:\n",
    "            img = tuple(img[i].detach() for i in range(len(img)))\n",
    "            if record_interval != -1:\n",
    "                img_list = tuple(Zip(*img_list, function=np.stack))\n",
    "        else:\n",
    "            img = img.detach()\n",
    "            if record_interval != -1:\n",
    "                img_list = np.stack(img_list)\n",
    "    if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "        neg_mask = tuple(neg_mask[i].detach() for i in range(model.mask_arity))\n",
    "        if record_interval != -1:\n",
    "            neg_mask_list = tuple(Zip(*neg_mask_list, function=np.stack))\n",
    "    if \"repr\" in args.ebm_target.split(\"+\"):\n",
    "        c_repr = c_repr.detach()\n",
    "        if record_interval != -1:\n",
    "            c_repr_list = np.stack(c_repr_list)\n",
    "    if \"z\" in args.ebm_target.split(\"+\"):\n",
    "        z = tuple(z[i].detach() if z[i] is not None else None for i in range(z_len))\n",
    "        if record_interval != -1:\n",
    "            z_list = tuple(Zip(*z_list, function=np.stack))\n",
    "    if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "        zgnn = tuple(zgnn[i].detach() if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "        if record_interval != -1:\n",
    "            zgnn_list = tuple(Zip(*zgnn_list))\n",
    "    if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "        wtarget = wtarget.detach()\n",
    "        if record_interval != -1:\n",
    "            wtarget_list = np.stack(wtarget_list)\n",
    "\n",
    "    info = {\n",
    "        \"neg_out_list\": neg_out_list,\n",
    "        \"img_list\": img_list,\n",
    "        \"neg_mask_list\": neg_mask_list,\n",
    "        \"c_repr_list\": c_repr_list,\n",
    "        \"z_list\": z_list,\n",
    "        \"zgnn_list\": zgnn_list,\n",
    "        \"wtarget_list\": wtarget_list,\n",
    "    }\n",
    "    reg_dict = transform_dict(reg_dict, \"array\")\n",
    "    info.update(reg_dict)\n",
    "    if is_return_E:\n",
    "        info[\"E_all\"] = E_all\n",
    "    return (img, neg_mask, c_repr, z, zgnn, wtarget), info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3.3 SGLD with KL:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def neg_mask_sgd_with_kl(\n",
    "    model,\n",
    "    img=None,\n",
    "    neg_mask=None,\n",
    "    c_repr=None,\n",
    "    z=None,\n",
    "    zgnn=None,\n",
    "    wtarget=None,\n",
    "    args=None,\n",
    "    mask_info=None,\n",
    "    is_return_E=False,\n",
    "    batch_shape=None,\n",
    "    record_interval=-1,\n",
    "):\n",
    "    \"\"\"Perform SGLD w.r.t. a subset of {img, neg_mask, c_repr, z} given the others, and in addition compute the the same thing for kl.\"\"\"\n",
    "    if args.step_size_img == -1:\n",
    "        args.step_size_img = args.step_size\n",
    "    if args.step_size_z == -1:\n",
    "        args.step_size_z = args.step_size\n",
    "    if args.step_size_zgnn == -1: \n",
    "        args.step_size_zgnn = args.step_size\n",
    "    if args.step_size_wtarget == -1:\n",
    "        args.step_size_wtarget = args.step_size\n",
    "    neg_out_list = []\n",
    "    # Initialize neg_mask and c_repr:\n",
    "    if img is not None:\n",
    "        batch_size = img[0].shape[0] if args.is_image_tuple else img.shape[0]\n",
    "        in_channels = img[0].shape[1] if args.is_image_tuple else img.shape[1]\n",
    "        device = img[0].device if args.is_image_tuple else img.device\n",
    "        img_shape = img[0].shape[2:] if args.is_image_tuple else img.shape[2:]\n",
    "    else:\n",
    "        batch_size = neg_mask[0].shape[0]\n",
    "        in_channels = args.in_channels\n",
    "        img_shape = args.image_size if args.rescaled_size == \"None\" else eval(args.rescaled_size)\n",
    "        device = neg_mask[0].device\n",
    "\n",
    "    img_list = []\n",
    "    if \"image\" in args.ebm_target.split(\"+\"):\n",
    "        img_value_min, img_value_max = args.image_value_range.split(\",\")\n",
    "        img_value_min, img_value_max = eval(img_value_min), eval(img_value_max)\n",
    "        img_value_span = img_value_max - img_value_min\n",
    "        assert img_value_span >= 1\n",
    "        if img is None:\n",
    "            img = (\n",
    "                torch.rand(batch_size, in_channels, *img_shape, device=device) * img_value_span + img_value_min, torch.rand(batch_size, in_channels, *img_shape, device=device) * img_value_span + img_value_min\n",
    "            ) if args.is_image_tuple else torch.rand(batch_size, in_channels, *img_shape, device=device) * img_value_span + img_value_min\n",
    "        else:\n",
    "            if isinstance(img, tuple):\n",
    "                img = tuple(torch.rand(batch_size, in_channels, *img_shape, device=device) * img_value_span + img_value_min if img_ele is None else img_ele for img_ele in img)\n",
    "        if record_interval != -1:\n",
    "            img_list.append(deepcopy(tuple(to_np_array(*img, keep_list=True)) if args.is_image_tuple else to_np_array(img)))\n",
    "\n",
    "    neg_mask_list = []\n",
    "    if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "        w_dim = 1 if \"mask\" in model.w_type else in_channels\n",
    "        if neg_mask is None:\n",
    "            neg_mask = tuple(torch.rand(batch_size, w_dim, *img_shape, device=device) for _ in range(model.mask_arity))\n",
    "        if record_interval != -1:\n",
    "            neg_mask_list.append(deepcopy(tuple(to_np_array(*neg_mask, keep_list=True))))\n",
    "\n",
    "    c_repr_list = []\n",
    "    if \"repr\" in args.ebm_target.split(\"+\"):\n",
    "        c_repr = torch.rand(batch_size, REPR_DIM, device=device) if c_repr is None else c_repr\n",
    "        if record_interval != -1:\n",
    "            c_repr_list.append(deepcopy(to_np_array(c_repr)))\n",
    "\n",
    "    z_list = []\n",
    "    if \"z\" in args.ebm_target.split(\"+\"):\n",
    "        z_len = 1\n",
    "        z = tuple(torch.rand(batch_size, model.z_dim, device=device) for _ in range(z_len)) if z is None else tuple(to_device_recur(z_ele, device) for z_ele in z)\n",
    "        z_len = len(z)\n",
    "        if record_interval != -1:\n",
    "            z_list.append(deepcopy(tuple(to_np_array(*z, keep_list=True))))\n",
    "    assert z is None or isinstance(z, tuple) or isinstance(z, list)\n",
    "\n",
    "    zgnn_list = []\n",
    "    if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "        n_nodes = len(z) if z is not None else len(neg_mask)\n",
    "        n_edges = n_nodes * (n_nodes - 1)\n",
    "        zgnn_dim = model.zgnn_dim\n",
    "        zgnn = (torch.rand(batch_shape[0], n_nodes, zgnn_dim, device=device) if model.gnn.is_zgnn_node else None, torch.rand(batch_shape[0], n_edges, model.edge_attr_size, device=device)) if zgnn is None else tuple(to_device_recur(zgnn_ele, device) for zgnn_ele in zgnn)\n",
    "        zgnn_len = len(zgnn)\n",
    "        if record_interval != -1:\n",
    "            zgnn_list.append(deepcopy(tuple(to_np_array(*zgnn, keep_list=True))))\n",
    "\n",
    "    wtarget_list = []\n",
    "    if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "        w_dim = 1 if \"mask\" in model.w_type else in_channels\n",
    "        if wtarget is None:\n",
    "            wtarget = torch.rand(batch_size, w_dim, *img_shape, device=device)\n",
    "        if record_interval != -1:\n",
    "            wtarget_list.append(deepcopy(to_np_array(wtarget)))\n",
    "\n",
    "    # Setting up noise and step_size scheduling:\n",
    "    if args.lambd_start == -1:\n",
    "        args.lambd_start = args.lambd\n",
    "    lambd_list = args.lambd + 1/2 * (args.lambd_start - args.lambd) * (1 + torch.cos(torch.arange(args.sample_step)/args.sample_step * np.pi))\n",
    "    if args.step_size_start == -1:\n",
    "        args.step_size_start = args.step_size\n",
    "    step_size_list = args.step_size + 1/2 * (args.step_size_start - args.step_size) * (1 + torch.cos(torch.arange(args.sample_step)/args.sample_step * np.pi))\n",
    "    if args.SGLD_is_anneal:\n",
    "        multiplier_list = np.linspace(0, 1, args.sample_step) ** args.SGLD_anneal_power\n",
    "    else:\n",
    "        multiplier_list = np.ones(args.sample_step)\n",
    "\n",
    "    if args.SGLD_object_exceed_coef > 0:\n",
    "        if isinstance(img, tuple):\n",
    "            pos_all_mask = ((img[0][:,:1] == 1).float(), (img[1][:,:1] == 1).float())\n",
    "        else:\n",
    "            pos_all_mask = (img[:,:1] == 1).float()\n",
    "\n",
    "    reg_dict = {}\n",
    "\n",
    "    # SGLD:\n",
    "    for k in range(args.sample_step):\n",
    "        multiplier = multiplier_list[k]\n",
    "        # Each step add noise using Langevin dynamics:\n",
    "        if \"image\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise_img\" not in locals():\n",
    "                noise_img = (torch.randn(batch_size, in_channels, *img_shape, device=device), torch.randn(batch_size, in_channels, *img_shape, device=device)) if args.is_image_tuple else torch.randn(batch_size, in_channels, *img_shape, device=device)\n",
    "            if lambd_list[k] > 0:\n",
    "                if args.is_image_tuple:\n",
    "                    for i in range(len(img)):\n",
    "                        noise_img[i].normal_(0, lambd_list[k])\n",
    "                    img = tuple(img[i] + noise_img[i] for i in range(len(img)))\n",
    "                else:\n",
    "                    noise_img.normal_(0, lambd_list[k])\n",
    "                    img = img + noise_img\n",
    "            if args.is_image_tuple:\n",
    "                for i in range(len(img)):\n",
    "                    img[i].requires_grad_(True)\n",
    "            else:\n",
    "                img.requires_grad_(True)\n",
    "        if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise\" not in locals():\n",
    "                w_dim = 1 if \"mask\" in model.w_type else in_channels\n",
    "                noise = tuple(torch.randn(batch_size, w_dim, *(args.image_size if args.rescaled_size == \"None\" else eval(args.rescaled_size)), device=device) for _ in range(model.mask_arity))\n",
    "            if lambd_list[k] > 0:\n",
    "                for i in range(model.mask_arity):\n",
    "                    noise[i].normal_(0, lambd_list[k])\n",
    "                neg_mask = tuple(neg_mask[i] + noise[i] for i in range(model.mask_arity))\n",
    "            for i in range(model.mask_arity):\n",
    "                neg_mask[i].requires_grad_(True)\n",
    "        if \"repr\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise_repr\" not in locals():\n",
    "                noise_repr = torch.randn(batch_size, REPR_DIM, device=device)\n",
    "            if lambd_list[k] > 0:\n",
    "                noise_repr.normal_(0, lambd_list[k])\n",
    "                c_repr = c_repr + noise_repr\n",
    "            c_repr.requires_grad_(True)\n",
    "        if \"z\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise_z\" not in locals():\n",
    "                noise_z = tuple(torch.randn(batch_size, model.z_dim, device=device) if z[i] is not None else None for i in range(z_len))\n",
    "            if lambd_list[k] > 0:\n",
    "                for i in range(z_len):\n",
    "                    if z[i] is not None:\n",
    "                        noise_z[i].normal_(0, lambd_list[k])\n",
    "                z = tuple(z[i] + noise_z[i] if z[i] is not None else None for i in range(z_len))\n",
    "            for i in range(z_len):\n",
    "                if z[i] is not None:\n",
    "                    z[i].requires_grad_(True)\n",
    "        if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise_zgnn\" not in locals():\n",
    "                noise_zgnn = tuple(torch.randn(zgnn[i].shape, device=device) if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "            if lambd_list[k] > 0:\n",
    "                for i in range(zgnn_len):\n",
    "                    if zgnn[i] is not None:\n",
    "                        noise_zgnn[i].normal_(0, lambd_list[k])\n",
    "                zgnn = tuple(zgnn[i] + noise_zgnn[i] if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "            for i in range(zgnn_len):\n",
    "                if zgnn[i] is not None:\n",
    "                    zgnn[i].requires_grad_(True)\n",
    "        if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "            if \"noise_wtarget\" not in locals():\n",
    "                w_dim = 1 if \"mask\" in model.w_type else in_channels\n",
    "                noise_wtarget = torch.randn(batch_size, w_dim, *(args.image_size if args.rescaled_size == \"None\" else eval(args.rescaled_size)), device=device)\n",
    "            if lambd_list[k] > 0:\n",
    "                noise_wtarget.normal_(0, lambd_list[k])\n",
    "                wtarget = wtarget + noise_wtarget\n",
    "            wtarget.requires_grad_(True)\n",
    "\n",
    "        # Compute neg_out and the gradient:\n",
    "        neg_out = model(img, neg_mask, c_repr=c_repr, z=z, zgnn=zgnn, wtarget=wtarget, batch_shape=batch_shape)\n",
    "        neg_out_sum = neg_out.sum()\n",
    "        record_data(reg_dict, deepcopy(to_np_array(neg_out)), \"neg_out_core_list\")\n",
    "\n",
    "        # Penalize areas where the masks overlap with each other\n",
    "        if args.SGLD_mutual_exclusive_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            mutual_exclusive = get_neg_mask_overlap(neg_mask, mask_info=mask_info, is_penalize_lower=args.SGLD_is_penalize_lower, img=img) * args.SGLD_mutual_exclusive_coef * multiplier\n",
    "            neg_out_sum += mutual_exclusive.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(mutual_exclusive)), \"mutual_exclusive_list\")\n",
    "        if args.SGLD_fine_mutual_exclusive_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            fine_mutual_exclusive = get_fine_neg_mask_overlap(neg_mask, mask_info=mask_info) * args.SGLD_fine_mutual_exclusive_coef * multiplier\n",
    "            neg_out_sum += fine_mutual_exclusive.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(fine_mutual_exclusive)), \"fine_mutual_exclusive_list\")\n",
    "        # Penalize areas where the negative mask hits areas not covered by the true mask\n",
    "        if args.SGLD_object_exceed_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            # SGLD_object_exceed_coef is annealed quadratically as the sample steps continue\n",
    "            object_exceed = get_neg_mask_exceed(neg_mask, pos_all_mask) * args.SGLD_object_exceed_coef * multiplier\n",
    "            neg_out_sum += object_exceed.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(object_exceed)), \"object_exceed_list\")\n",
    "        # Encourage each EBM will either fully explain or fully ignore each pixel:\n",
    "        if args.SGLD_pixel_entropy_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            pixel_entropy = get_pixel_entropy(neg_mask) * args.SGLD_pixel_entropy_coef * multiplier\n",
    "            neg_out_sum += pixel_entropy.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(pixel_entropy)), \"pixel_entropy_list\")\n",
    "        # Encourage each EBM mask will either fully explain an object or not explain any:\n",
    "        if args.SGLD_mask_entropy_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            mask_entropy = get_mask_entropy(neg_mask) * args.SGLD_mask_entropy_coef * multiplier\n",
    "            neg_out_sum += mask_entropy.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(mask_entropy)), \"mask_entropy_list\")\n",
    "        # Encourage the ebms to specialize, where the ebms whose masks that are nearest to 0 or 1 will be push harder towards 0 or 1:\n",
    "        if args.SGLD_pixel_gm_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            pixel_gm = get_pixel_gm(neg_mask) * args.SGLD_pixel_gm_coef * multiplier\n",
    "            neg_out_sum += pixel_gm.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(pixel_gm)), \"pixel_gm_list\")\n",
    "        # Compute the energies that encourages selector discovery:\n",
    "        if \"mask\" in args.ebm_target.split(\"+\") and (\n",
    "            args.SGLD_iou_batch_consistency_coef > 0 or\n",
    "            args.SGLD_iou_concept_repel_coef > 0 or\n",
    "            args.SGLD_iou_relation_repel_coef > 0 or\n",
    "            args.SGLD_iou_relation_overlap_coef > 0 or\n",
    "            args.SGLD_iou_attract_coef > 0):\n",
    "            graph_energy = get_graph_energy(\n",
    "                neg_mask,\n",
    "                mask_info=mask_info,\n",
    "                iou_batch_consistency_coef=args.SGLD_iou_batch_consistency_coef,\n",
    "                iou_concept_repel_coef=args.SGLD_iou_concept_repel_coef,\n",
    "                iou_relation_repel_coef=args.SGLD_iou_relation_repel_coef,\n",
    "                iou_relation_overlap_coef=args.SGLD_iou_relation_overlap_coef,\n",
    "                iou_attract_coef=args.SGLD_iou_attract_coef,\n",
    "                batch_shape=batch_shape,\n",
    "            )[0] * multiplier\n",
    "            neg_out_sum += graph_energy.sum()\n",
    "            record_data(reg_dict, deepcopy(to_np_array(graph_energy)), \"graph_energy_list\")\n",
    "\n",
    "        if \"image\" in args.ebm_target.split(\"+\"):\n",
    "            if args.is_image_tuple:\n",
    "                img_grad = tuple(torch.autograd.grad([neg_out_sum],\n",
    "                                  [img[i]],\n",
    "                                  create_graph=True if args.kl_all_step else False,\n",
    "                                  retain_graph=True if \"mask\" in args.ebm_target.split(\"+\") or \"repr\" in args.ebm_target.split(\"+\") or \"z\" in args.ebm_target.split(\"+\") or \"zgnn\" in args.ebm_target.split(\"+\") or \"wtarget\" in args.ebm_target.split(\"+\") or i < len(img) - 1 else None)[0] for i in range(len(img)))\n",
    "            else:\n",
    "                img_grad = torch.autograd.grad([neg_out_sum],\n",
    "                                               [img],\n",
    "                                               create_graph=True if args.kl_all_step else False,\n",
    "                                               retain_graph=True if \"mask\" in args.ebm_target.split(\"+\") or \"repr\" in args.ebm_target.split(\"+\") or \"z\" in args.ebm_target.split(\"+\") or \"zgnn\" in args.ebm_target.split(\"+\") or \"wtarget\" in args.ebm_target.split(\"+\") else None)[0]\n",
    "        if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            neg_mask_grad = tuple(torch.autograd.grad([neg_out_sum],\n",
    "                                  [neg_mask[i]],\n",
    "                                  create_graph=True if args.kl_all_step else False,\n",
    "                                  retain_graph=True if \"repr\" in args.ebm_target.split(\"+\") or \"z\" in args.ebm_target.split(\"+\") or \"zgnn\" in args.ebm_target.split(\"+\") or \"wtarget\" in args.ebm_target.split(\"+\") or i < model.mask_arity - 1 else None)[0] for i in range(model.mask_arity))\n",
    "        if \"repr\" in args.ebm_target.split(\"+\"):\n",
    "            c_repr_grad = torch.autograd.grad([neg_out_sum], [c_repr], create_graph=True if args.kl_all_step else False,\n",
    "                                              retain_graph=True if \"z\" in args.ebm_target.split(\"+\") or \"zgnn\" in args.ebm_target.split(\"+\") or \"wtarget\" in args.ebm_target.split(\"+\") else None\n",
    "                                             )[0]\n",
    "        if \"z\" in args.ebm_target.split(\"+\"):\n",
    "            z_grad = tuple(torch.autograd.grad([neg_out_sum],\n",
    "                                               [z[i]],\n",
    "                                               create_graph=True if args.kl_all_step else False,\n",
    "                                               retain_graph=True if \"zgnn\" in args.ebm_target.split(\"+\") or \"wtarget\" in args.ebm_target.split(\"+\") or i < z_len - 1 else None)[0] if z[i] is not None else None for i in range(z_len))\n",
    "        if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "            zgnn_grad = tuple(torch.autograd.grad([neg_out_sum],\n",
    "                                                  [zgnn[i]],\n",
    "                                                  create_graph=True if args.kl_all_step else False,\n",
    "                                                  retain_graph=True if \"wtarget\" in args.ebm_target.split(\"+\") or i < zgnn_len - 1 else None)[0] if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "        if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "            wtarget_grad = torch.autograd.grad([neg_out_sum], [wtarget], create_graph=True if args.kl_all_step else False)[0]\n",
    "\n",
    "        if k == args.sample_step - 1:\n",
    "            img_ori = img\n",
    "            neg_mask_ori = neg_mask\n",
    "            c_repr_ori = c_repr\n",
    "            z_ori = z\n",
    "            zgnn_ori = zgnn\n",
    "            wtarget_ori = wtarget\n",
    "            # Update subset of {img, neg_mask, c_repr, z} using the gradient:\n",
    "            if \"image\" in args.ebm_target.split(\"+\"):\n",
    "                if args.is_image_tuple:\n",
    "                    img = tuple(img[i] - img_grad[i] * args.step_size_img for i in range(len(img)))\n",
    "\n",
    "                    # Update the img_kl:\n",
    "                    img_kl = img_ori\n",
    "                    neg_out_kl_sum = model(img_kl, neg_mask, c_repr=c_repr, z=z, zgnn=zgnn, wtarget=wtarget, batch_shape=batch_shape).sum()\n",
    "                    img_kl_grad = tuple(torch.autograd.grad([neg_out_kl_sum], [img_kl[i]], create_graph=True)[0] for i in range(len(img)))\n",
    "                    img_kl = tuple(img_kl[i] - img_kl_grad[i] * args.step_size_img for i in range(len(img)))\n",
    "                    img_kl = tuple(torch.clamp(img_kl[i], img_value_min, img_value_max) for i in range(len(img)))\n",
    "\n",
    "                    # Detach and clamp neg_mask:\n",
    "                    img = tuple(img[i].detach() for i in range(len(img)))\n",
    "                    img = tuple(torch.clamp(img[i], img_value_min, img_value_max) for i in range(len(img)))\n",
    "                    if record_interval != -1:\n",
    "                        img_list.append(deepcopy(tuple(to_np_array(*img, keep_list=True))))\n",
    "                else:\n",
    "                    img = img - img_grad * args.step_size_img\n",
    "\n",
    "                    # Update the img_kl:\n",
    "                    img_kl = img_ori\n",
    "                    neg_out_kl_sum = model(img_kl, neg_mask, c_repr=c_repr, z=z, zgnn=zgnn, wtarget=wtarget, batch_shape=batch_shape).sum()\n",
    "                    img_kl_grad = torch.autograd.grad([neg_out_kl_sum], [img_kl], create_graph=True)[0]\n",
    "                    img_kl = img_kl - img_kl_grad * args.step_size_img\n",
    "                    img_kl = torch.clamp(img_kl, img_value_min, img_value_max)\n",
    "\n",
    "                    # Detach and clamp z:\n",
    "                    img = img.detach()\n",
    "                    img = torch.clamp(img, img_value_min, img_value_max)\n",
    "                    if record_interval != -1:\n",
    "                        img_list.append(deepcopy(to_np_array(img)))\n",
    "\n",
    "            if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "                neg_mask = tuple(neg_mask[i] - neg_mask_grad[i] * step_size_list[k] for i in range(model.mask_arity))\n",
    "\n",
    "                # Update the neg_mask_kl:\n",
    "                neg_mask_kl = neg_mask_ori\n",
    "                if is_return_E:\n",
    "                    neg_out_kl_sum = model(img_ori, neg_mask_kl, c_repr=c_repr, z=z, zgnn=zgnn, wtarget=wtarget, batch_shape=batch_shape, is_E_tensor=is_return_E).sum()\n",
    "                    E_all = model.info[\"E_all\"]\n",
    "                else:\n",
    "                    neg_out_kl_sum = model(img_ori, neg_mask_kl, c_repr=c_repr, z=z, zgnn=zgnn, wtarget=wtarget, batch_shape=batch_shape).sum()\n",
    "\n",
    "                if args.SGLD_mutual_exclusive_coef > 0:\n",
    "                    neg_out_kl_sum += get_neg_mask_overlap(neg_mask_kl, mask_info=mask_info, is_penalize_lower=args.SGLD_is_penalize_lower, img=img).sum() * args.SGLD_mutual_exclusive_coef * multiplier\n",
    "                if args.SGLD_fine_mutual_exclusive_coef > 0 and \"mask\" in args.ebm_target.split(\"+\"):\n",
    "                    neg_out_kl_sum += get_fine_neg_mask_overlap(neg_mask_kl, mask_info=mask_info).sum() * args.SGLD_fine_mutual_exclusive_coef * multiplier\n",
    "                # Penalize areas where the negative mask hits areas not covered by the true mask\n",
    "                if args.SGLD_object_exceed_coef > 0:\n",
    "                    # SGLD_object_exceed_coef is annealed quadratically as the sample steps continue\n",
    "                    neg_out_kl_sum += get_neg_mask_exceed(neg_mask_kl, pos_all_mask).sum() * args.SGLD_object_exceed_coef * multiplier\n",
    "                # Encourage each EBM will either fully explain or fully ignore each pixel:\n",
    "                if args.SGLD_pixel_entropy_coef > 0:\n",
    "                    neg_out_kl_sum += get_pixel_entropy(neg_mask_kl).sum() * args.SGLD_pixel_entropy_coef * multiplier\n",
    "                # Encourage each EBM mask will either fully explain an object or not explain any:\n",
    "                if args.SGLD_mask_entropy_coef > 0:\n",
    "                    neg_out_kl_sum += get_mask_entropy(neg_mask_kl).sum() * args.SGLD_mask_entropy_coef * multiplier\n",
    "                # Encourage the ebms to specialize, where the ebms whose masks that are nearest to 0 or 1 will be push harder towards 0 or 1:\n",
    "                if args.SGLD_pixel_gm_coef > 0:\n",
    "                    neg_out_kl_sum += get_pixel_gm(neg_mask_kl).sum() * args.SGLD_pixel_gm_coef * multiplier\n",
    "                # Compute the energies that encourages selector discovery:\n",
    "                if \"mask\" in args.ebm_target.split(\"+\") and (\n",
    "                    args.SGLD_iou_batch_consistency_coef > 0 or\n",
    "                    args.SGLD_iou_concept_repel_coef > 0 or\n",
    "                    args.SGLD_iou_relation_repel_coef > 0 or\n",
    "                    args.SGLD_iou_relation_overlap_coef > 0 or\n",
    "                    args.SGLD_iou_attract_coef > 0):\n",
    "                    neg_out_kl_sum += get_graph_energy(\n",
    "                        neg_mask_kl,\n",
    "                        mask_info=mask_info,\n",
    "                        iou_batch_consistency_coef=args.SGLD_iou_batch_consistency_coef,\n",
    "                        iou_concept_repel_coef=args.SGLD_iou_concept_repel_coef,\n",
    "                        iou_relation_repel_coef=args.SGLD_iou_relation_repel_coef,\n",
    "                        iou_relation_overlap_coef=args.SGLD_iou_relation_overlap_coef,\n",
    "                        iou_attract_coef=args.SGLD_iou_attract_coef,\n",
    "                        batch_shape=batch_shape,\n",
    "                    )[0].sum() * multiplier\n",
    "\n",
    "                neg_mask_kl_grad = tuple(torch.autograd.grad([neg_out_kl_sum], [neg_mask_kl[i]], create_graph=True)[0] for i in range(model.mask_arity))\n",
    "                neg_mask_kl = tuple(neg_mask_kl[i] - neg_mask_kl_grad[i] * step_size_list[k] for i in range(model.mask_arity))\n",
    "                neg_mask_kl = tuple(torch.clamp(neg_mask_kl[i], 0, 1) for i in range(model.mask_arity))\n",
    "\n",
    "                # Detach and clamp neg_mask:\n",
    "                neg_mask = tuple(neg_mask[i].detach() for i in range(model.mask_arity))\n",
    "                neg_mask = tuple(torch.clamp(neg_mask[i], 0, 1) for i in range(model.mask_arity))\n",
    "                if record_interval != -1:\n",
    "                    neg_mask_list.append(deepcopy(tuple(to_np_array(*neg_mask, keep_list=True))))\n",
    "            if \"repr\" in args.ebm_target.split(\"+\"):\n",
    "                c_repr = c_repr - c_repr_grad * args.step_size_repr\n",
    "\n",
    "                # Update the c_repr_kl:\n",
    "                c_repr_kl = c_repr_ori\n",
    "                neg_out_kl_sum = model(img_ori, neg_mask_ori, c_repr=c_repr_kl, z=z, zgnn=zgnn, wtarget=wtarget, batch_shape=batch_shape).sum()\n",
    "                # No need for SGLD_mutual_exclusive_coef since the grad is w.r.t. c_repr_kl.\n",
    "                c_repr_kl_grad = torch.autograd.grad([neg_out_kl_sum], [c_repr_kl], create_graph=True)[0]\n",
    "                c_repr_kl = c_repr_kl - c_repr_kl_grad * args.step_size_repr\n",
    "                if \"softmax\" not in args.c_repr_mode:\n",
    "                    c_repr_kl = torch.clamp(c_repr_kl, 0, 1)\n",
    "\n",
    "                # Detach and clamp c_repr:\n",
    "                c_repr = c_repr.detach()\n",
    "                if \"softmax\" not in args.c_repr_mode:\n",
    "                    c_repr = torch.clamp(c_repr, 0, 1)\n",
    "                if record_interval != -1:\n",
    "                    c_repr_list.append(deepcopy(to_np_array(c_repr)))\n",
    "            if \"z\" in args.ebm_target.split(\"+\"):\n",
    "                z = tuple(z[i] - z_grad[i] * args.step_size_z if z[i] is not None else None for i in range(z_len))\n",
    "\n",
    "                # Update the z_kl:\n",
    "                z_kl = z_ori\n",
    "                neg_out_kl_sum = model(img_ori, neg_mask_ori, c_repr=c_repr_ori, z=z_kl, zgnn=zgnn, wtarget=wtarget, batch_shape=batch_shape).sum()\n",
    "                z_kl_grad = tuple(torch.autograd.grad([neg_out_kl_sum], [z_kl[i]], create_graph=True)[0] if z[i] is not None else None for i in range(z_len))\n",
    "                z_kl = tuple(z_kl[i] - z_kl_grad[i] * args.step_size_z if z[i] is not None else None for i in range(z_len))\n",
    "                z_kl = tuple(torch.clamp(z_kl[i], 0, 1) if z[i] is not None else None for i in range(z_len))\n",
    "\n",
    "                # Detach and clamp z:\n",
    "                z = tuple(z[i].detach() if z[i] is not None else None for i in range(z_len))\n",
    "                z = tuple(torch.clamp(z[i], 0, 1) if z[i] is not None else None for i in range(z_len))\n",
    "                if record_interval != -1:\n",
    "                    z_list.append(deepcopy(tuple(to_np_array(*z, keep_list=True))))\n",
    "            if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "                zgnn = tuple(zgnn[i] - zgnn_grad[i] * args.step_size_zgnn if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "\n",
    "                # Update the zgnn_kl:\n",
    "                zgnn_kl = zgnn_ori\n",
    "                neg_out_kl_sum = model(img_ori, neg_mask_ori, c_repr=c_repr_ori, z=z_ori, zgnn=zgnn_kl, wtarget=wtarget, batch_shape=batch_shape).sum()\n",
    "                zgnn_kl_grad = tuple(torch.autograd.grad([neg_out_kl_sum], [zgnn_kl[i]], create_graph=True)[0] if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "                zgnn_kl = tuple(zgnn_kl[i] - zgnn_kl_grad[i] * args.step_size_zgnn if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "                zgnn_kl = tuple(torch.clamp(zgnn_kl[i], 0, 1) if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "\n",
    "                # Detach and clamp z:\n",
    "                zgnn = tuple(zgnn[i].detach() if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "                zgnn = tuple(torch.clamp(zgnn[i], 0, 1) if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "                if record_interval != -1:\n",
    "                    zgnn_list.append(deepcopy(tuple(to_np_array(*zgnn, keep_list=True))))\n",
    "            if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "                wtarget = wtarget - wtarget_grad * args.step_size_wtarget\n",
    "\n",
    "                # Update the wtarget_kl:\n",
    "                wtarget_kl = wtarget_ori\n",
    "                neg_out_kl_sum = model(img_ori, neg_mask_ori, c_repr=c_repr_ori, z=z_ori, zgnn=zgnn_ori, wtarget=wtarget_kl, batch_shape=batch_shape).sum()\n",
    "                wtarget_kl_grad = torch.autograd.grad([neg_out_kl_sum], [wtarget_kl], create_graph=True)[0]\n",
    "                wtarget_kl = wtarget_kl - wtarget_kl_grad * args.step_size_wtarget\n",
    "                wtarget_kl = torch.clamp(wtarget_kl, 0, 1)\n",
    "\n",
    "                # Detach and clamp z:\n",
    "                wtarget = wtarget.detach()\n",
    "                wtarget = torch.clamp(wtarget, 0, 1)\n",
    "                if record_interval != -1:\n",
    "                    wtarget_list.append(deepcopy(to_np_array(wtarget)))\n",
    "\n",
    "        else:\n",
    "            # Update subset of {img, neg_mask, c_repr, z} using the gradient:\n",
    "            if \"image\" in args.ebm_target.split(\"+\"):\n",
    "                if args.is_image_tuple:\n",
    "                    img = tuple(img[i] - img_grad[i] * args.step_size_img for i in range(len(img)))\n",
    "                    img = tuple(img[i].detach() for i in range(len(img)))\n",
    "                    img = tuple(torch.clamp(img[i], img_value_min, img_value_max) for i in range(len(img)))\n",
    "                    if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                        img_list.append(deepcopy(tuple(to_np_array(*img, keep_list=True))))\n",
    "                else:\n",
    "                    img = img - img_grad * args.step_size_img\n",
    "                    img = img.detach()\n",
    "                    img = torch.clamp(img, img_value_min, img_value_max)\n",
    "                    if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                        img_list.append(deepcopy(to_np_array(img)))\n",
    "            if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "                neg_mask = tuple(neg_mask[i] - neg_mask_grad[i]*step_size_list[k] for i in range(model.mask_arity))\n",
    "                neg_mask = tuple(neg_mask[i].detach() for i in range(model.mask_arity))\n",
    "                neg_mask = tuple(torch.clamp(neg_mask[i], 0, 1) for i in range(model.mask_arity))\n",
    "                if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                    neg_mask_list.append(deepcopy(tuple(to_np_array(*neg_mask, keep_list=True))))\n",
    "            if \"repr\" in args.ebm_target.split(\"+\"):\n",
    "                c_repr = c_repr - c_repr_grad * args.step_size_repr\n",
    "                c_repr = c_repr.detach()\n",
    "                if \"softmax\" not in args.c_repr_mode:\n",
    "                    c_repr = torch.clamp(c_repr, 0, 1)\n",
    "                if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                    c_repr_list.append(deepcopy(to_np_array(c_repr)))\n",
    "            if \"z\" in args.ebm_target.split(\"+\"):\n",
    "                z = tuple(z[i] - z_grad[i] * args.step_size_z if z[i] is not None else None for i in range(z_len))\n",
    "                z = tuple(z[i].detach() if z[i] is not None else None for i in range(z_len))\n",
    "                z = tuple(torch.clamp(z[i], 0, 1) if z[i] is not None else None for i in range(z_len))\n",
    "                if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                    z_list.append(deepcopy(tuple(to_np_array(*z, keep_list=True))))\n",
    "            if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "                zgnn = tuple(zgnn[i] - zgnn_grad[i] * args.step_size_zgnn if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "                zgnn = tuple(zgnn[i].detach() if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "                zgnn = tuple(torch.clamp(zgnn[i], 0, 1) if zgnn[i] is not None else None for i in range(zgnn_len))\n",
    "                if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                    zgnn_list.append(deepcopy(tuple(to_np_array(*zgnn, keep_list=True))))\n",
    "            if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "                wtarget = wtarget - wtarget_grad * args.step_size_wtarget\n",
    "                wtarget = wtarget.detach()\n",
    "                wtarget = torch.clamp(wtarget, 0, 1)\n",
    "                if record_interval != -1 and ((k + 1) % record_interval == 0 or k == args.sample_step - 1):\n",
    "                    wtarget_list.append(deepcopy(to_np_array(wtarget)))\n",
    "        neg_out_list.append(deepcopy(to_np_array(neg_out)))\n",
    "\n",
    "    neg_out_list = np.concatenate(neg_out_list, -1).T   # after: [sample_step, B]\n",
    "    if record_interval != -1:\n",
    "        if \"image\" in args.ebm_target.split(\"+\"):\n",
    "            if args.is_image_tuple:\n",
    "                img_list = tuple(Zip(*img_list, function=np.stack))\n",
    "            else:\n",
    "                img_list = np.stack(img_list)\n",
    "        if \"mask\" in args.ebm_target.split(\"+\"):\n",
    "            neg_mask_list = tuple(Zip(*neg_mask_list, function=np.stack))\n",
    "        if \"repr\" in args.ebm_target.split(\"+\"):\n",
    "            c_repr_list = np.stack(c_repr_list)\n",
    "        if \"z\" in args.ebm_target.split(\"+\"):\n",
    "            z_list = tuple(Zip(*z_list, function=np.stack))\n",
    "        if \"zgnn\" in args.ebm_target.split(\"+\"):\n",
    "            zgnn_list = tuple(Zip(*zgnn_list))\n",
    "        if \"wtarget\" in args.ebm_target.split(\"+\"):\n",
    "            wtarget_list = np.stack(wtarget)\n",
    "\n",
    "    info = {\n",
    "        \"neg_out_list\": neg_out_list,\n",
    "        \"img_list\": img_list,\n",
    "        \"neg_mask_list\": neg_mask_list,\n",
    "        \"c_repr_list\": c_repr_list,\n",
    "        \"z_list\": z_list,\n",
    "        \"zgnn_list\": zgnn_list,\n",
    "        \"wtarget_list\": wtarget_list,\n",
    "    }\n",
    "    if \"img_kl\" not in locals():\n",
    "        img_kl = None\n",
    "    if \"neg_mask_kl\" not in locals():\n",
    "        neg_mask_kl = None\n",
    "    if \"c_repr_kl\" not in locals():\n",
    "        c_repr_kl = None\n",
    "    if \"z_kl\" not in locals():\n",
    "        z_kl = None\n",
    "    if \"zgnn_kl\" not in locals():\n",
    "        zgnn_kl = None\n",
    "    if \"wtarget_kl\" not in locals():\n",
    "        wtarget_kl = None\n",
    "    reg_dict = transform_dict(reg_dict, \"array\")\n",
    "    info.update(reg_dict)\n",
    "    if is_return_E:\n",
    "        info[\"E_all\"] = E_all\n",
    "    return (img, neg_mask, c_repr, z, zgnn, wtarget), (img_kl, neg_mask_kl, c_repr_kl, z_kl, zgnn_kl, wtarget_kl), info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3.4 Helper functions for objectives:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_graph_energy(\n",
    "    mask,\n",
    "    mask_info,\n",
    "    iou_batch_consistency_coef=0,\n",
    "    iou_concept_repel_coef=0,\n",
    "    iou_relation_repel_coef=0,\n",
    "    iou_relation_overlap_coef=0,\n",
    "    iou_attract_coef=0,\n",
    "    batch_shape=None,\n",
    "):\n",
    "    \"\"\"\n",
    "    Obtain the three graph energy that encourages forming a common graph among several examples.\n",
    "\n",
    "    Args:\n",
    "        mask: a tuple of masks, each with shape [B, 1, H, W]\n",
    "        mask_info: a dictionary containing information about the masks. E.g.\n",
    "            {\n",
    "                id_to_type: {0: (\"concept\", 1), 1: (\"relation\", 0), 2: (\"concept\", 0), ...},  \n",
    "                    # The number are chosen from {0,1}, which indicates the number of object slot this mask occupies, for computing mutual_exclusive loss.\n",
    "                id_same_relation: [(1,3), (4,6), ...],\n",
    "            }\n",
    "        batch_shape: if not None, will have shape of [B_task, B_example]\n",
    "    \"\"\"\n",
    "    def get_triu_ids(array, is_triu=True):\n",
    "        if isinstance(array, Number):\n",
    "            array = np.arange(array)\n",
    "        rows_matrix, col_matrix = np.meshgrid(array, array)\n",
    "        matrix_cat = np.stack([rows_matrix, col_matrix], -1)\n",
    "        rr, cc = np.triu_indices(len(matrix_cat), k=1)\n",
    "        rows, cols = matrix_cat[cc, rr].T\n",
    "        return rows, cols\n",
    "\n",
    "    def get_relation_overlap_loss_pair(distance_matrix_batch, rel_tuple_1, rel_tuple_2):\n",
    "        \"\"\"distance_matrix_batch: [B_task, B_example, n_masks, n_masks]\"\"\"\n",
    "        shape = distance_matrix_batch.shape\n",
    "        assert len(shape) == 4\n",
    "        assert shape[2] == shape[3]\n",
    "        rel_tuple_2_r = (rel_tuple_2[1], rel_tuple_2[0])\n",
    "        loss = (1 - distance_matrix_batch[:, :, rel_tuple_1, rel_tuple_2]).clamp(0, 1).prod(-1) + \\\n",
    "               (1 - distance_matrix_batch[:, :, rel_tuple_1, rel_tuple_2_r]).clamp(0, 1).prod(-1)\n",
    "        return loss\n",
    "\n",
    "    assert isinstance(mask, tuple) or isinstance(mask, list)\n",
    "    assert mask[0].shape[1] == 1\n",
    "    n_masks = len(mask)\n",
    "    device = mask[0].device\n",
    "    batch_size = mask[0].shape[0]\n",
    "    loss_dict = {}\n",
    "\n",
    "    # Get pairwise Jaccard distance between masks:\n",
    "    mask = torch.stack(mask, 1)  # [B, n_masks, 1, H, W]\n",
    "    if batch_shape is None:\n",
    "        mask = mask[None]  # [B_task:1, B_example, n_masks, 1, H, W]\n",
    "        batch_shape = (1, batch_size)\n",
    "    else:\n",
    "        mask = mask.view(*batch_shape, *mask.shape[1:])  # [B_task, B_example, n_masks, 1, H, W]\n",
    "    mask_expand_0 = mask[:,:,:,None]\n",
    "    mask_expand_1 = mask[:,:,None]\n",
    "    distance_matrix_batch = get_soft_Jaccard_distance(mask_expand_0, mask_expand_1, dim=(-3,-2,-1))  # [B_task, B_example, n_masks, n_masks]\n",
    "    distance_matrix_mean = distance_matrix_batch.mean(1, keepdims=True)  # [B_task, 1, n_masks, n_masks]\n",
    "\n",
    "    # Encourage that the pairs in the same batch has similar distances across examples:\n",
    "    loss_all = torch.FloatTensor([[[0]]]).to(device)\n",
    "    if iou_batch_consistency_coef > 0 and batch_size > 1:\n",
    "        loss_batch_consistency = ((distance_matrix_batch - distance_matrix_mean).square().mean(1) + 1e-10).sqrt().mean((1,2), keepdims=True).expand(-1, batch_shape[1], 1) * iou_batch_consistency_coef  # [B_task, B_example, 1]\n",
    "        loss_all = loss_all + loss_batch_consistency\n",
    "        loss_dict[\"loss_batch_consistency\"] = to_np_array(loss_batch_consistency)\n",
    "\n",
    "    # Encourage that the masks for different concepts do not overlap:\n",
    "    repel_rows = []\n",
    "    repel_cols = []\n",
    "    concept_ids = [id for id, item in mask_info[\"id_to_type\"].items() if item[0] == \"concept\"]\n",
    "    concept_rows, concept_cols = get_triu_ids(concept_ids)\n",
    "    repel_rows.append(concept_rows)\n",
    "    repel_cols.append(concept_cols)\n",
    "    if iou_concept_repel_coef > 0:\n",
    "        loss_concept_repel = (1 - distance_matrix_batch[:,:,concept_rows,concept_cols].mean(-1, keepdims=True)) * iou_concept_repel_coef  # [B_task, B_example, 1]\n",
    "        loss_all = loss_all + loss_concept_repel\n",
    "        loss_dict[\"loss_concept_repel\"] = to_np_array(loss_concept_repel)\n",
    "\n",
    "    # Encourage the mask for the same relation do not overlap:\n",
    "    if len(mask_info[\"id_same_relation\"]) > 0:\n",
    "        relation_rows, relation_cols = np.stack(mask_info[\"id_same_relation\"]).T\n",
    "    else:\n",
    "        relation_rows, relation_cols = np.array([]), np.array([])\n",
    "    repel_rows.append(relation_rows)\n",
    "    repel_cols.append(relation_cols)\n",
    "    if iou_relation_repel_coef > 0 and len(mask_info[\"id_same_relation\"]) > 0:\n",
    "        loss_relation_repel = (1 - distance_matrix_batch[:,:,relation_rows,relation_cols].mean(-1, keepdims=True)) * iou_relation_repel_coef  # [B_task, B_example, 1]\n",
    "        loss_all = loss_all + loss_relation_repel\n",
    "        loss_dict[\"loss_relation_repel\"] = to_np_array(loss_relation_repel)\n",
    "\n",
    "    # Discourage two relation-EBMs to discover the same pair of objects:\n",
    "    if iou_relation_overlap_coef > 0 and len(mask_info[\"id_same_relation\"]) > 1:\n",
    "        id_same_relation = mask_info[\"id_same_relation\"]\n",
    "        length = len(id_same_relation)\n",
    "        pairs = [(id_same_relation[i], id_same_relation[j]) for i in range(length) for j in range(length) if i < j]\n",
    "        loss_list = [get_relation_overlap_loss_pair(distance_matrix_batch, rel_tuple_1, rel_tuple_2) for rel_tuple_1, rel_tuple_2 in pairs]\n",
    "        loss_relation_overlap = torch.stack(loss_list, -1).sum(-1, keepdims=True) * iou_relation_overlap_coef\n",
    "        loss_all = loss_all + loss_relation_overlap\n",
    "        loss_dict[\"loss_relation_overlap\"] = to_np_array(loss_relation_overlap)\n",
    "\n",
    "    # Encourage that compatible (not repelled) masks can snap to each other when they are near enough:\n",
    "    if iou_attract_coef > 0:\n",
    "        all_rows, all_cols = get_triu_ids(n_masks)\n",
    "        repel_rows = np.concatenate(repel_rows)\n",
    "        repel_cols = np.concatenate(repel_cols)\n",
    "        all_tuples = [(row, col) for row, col in zip(all_rows, all_cols)]\n",
    "        repel_tuples = [(row, col) for row, col in zip(repel_rows, repel_cols)]\n",
    "        compat_tuples = [ele for ele in all_tuples if ele not in repel_tuples]\n",
    "        if len(compat_tuples) > 0:\n",
    "            compat_rows, compat_cols = np.array(compat_tuples).T\n",
    "            distance_attract = distance_matrix_mean[:,:,compat_rows, compat_cols].clamp(1e-5)  # [B_task, 1, n_compat]\n",
    "            print(\"\\nDistance_attract:\")\n",
    "            print(distance_attract)\n",
    "            loss_attract = distance_attract.mean(-1, keepdims=True).expand(-1, batch_shape[1], 1) * iou_attract_coef  # [B_task, B_example, 1]\n",
    "            loss_all = loss_all + loss_attract\n",
    "            loss_dict[\"loss_attract\"] = to_np_array(loss_attract)\n",
    "    return loss_all, loss_dict\n",
    "\n",
    "\n",
    "def get_pixel_gm(mask_list, order=-1, epsilon=1e-5):\n",
    "    \"\"\"Get generalized-mean over the distance to 0 and 1 over multiple masks for each pixel.\n",
    "        Use an order <=0 or \"min\" to encourge specialization of masks.\n",
    "    \"\"\"\n",
    "    assert isinstance(mask_list, list) or isinstance(mask_list, tuple)\n",
    "    shape = mask_list[0].shape\n",
    "    assert len(shape) == 4 and shape[1] == 1\n",
    "    n_masks = len(mask_list)\n",
    "\n",
    "    mask_list = torch.stack(mask_list, -1)  # Last dimension is the n_masks dimension\n",
    "    if isinstance(order, Number):\n",
    "        mask_list = mask_list.clamp(epsilon, 1-epsilon)\n",
    "\n",
    "    if order == -1:\n",
    "        L1 = (n_masks / (1 / mask_list).sum(-1)).mean((1,2,3))\n",
    "        L2 = (n_masks / (1 / (1-mask_list)).sum(-1)).mean((1,2,3))\n",
    "    elif order == 0:\n",
    "        L1 = (mask_list.prod(-1) ** (1/n_masks)).mean((1,2,3))\n",
    "        L2 = ((1-mask_list).prod(-1) ** (1/n_masks)).mean((1,2,3))\n",
    "    elif order == 1:\n",
    "        L1 = mask_list.mean((1,2,3,4))\n",
    "        L2 = (1-mask_list).mean((1,2,3,4))\n",
    "    elif order == \"max\":\n",
    "        L1 = mask_list.max(4)[0].mean((1,2,3))\n",
    "        L2 = (1-mask_list).max(-1)[0].mean((1,2,3))\n",
    "    elif order == \"min\":\n",
    "        L1 = mask_list.min(4)[0].mean((1,2,3))\n",
    "        L2 = (1-mask_list).min(-1)[0].mean((1,2,3))\n",
    "    else:\n",
    "        assert isinstance(order, Number)\n",
    "        L1 = (((mask_list ** order).mean(-1)) ** (1 / float(order))).mean((1,2,3))\n",
    "        L2 = ((((1-mask_list) ** order).mean(-1)) ** (1 / float(order))).mean((1,2,3))\n",
    "    pixel_gm = L1 + L2\n",
    "    return pixel_gm\n",
    "\n",
    "\n",
    "def get_pixel_entropy(mask_list, epsilon=1e-5):\n",
    "    \"\"\"Obtain the pixel-wise entropy for each mask.\n",
    "    Args:\n",
    "        mask_list: each mask in mask_list should have the shape of [B, 1, H, W]\n",
    "        epsilon: to prevent NaN.\n",
    "\n",
    "    Returns:\n",
    "        entropy: shape [B,]\n",
    "    \"\"\"\n",
    "    # Make sure that the format is correct:\n",
    "    assert isinstance(mask_list, list) or isinstance(mask_list, tuple)\n",
    "    shape = mask_list[0].shape\n",
    "    assert len(shape) == 4 and shape[1] == 1\n",
    "    \n",
    "    mask_list = torch.stack(mask_list, 1).clamp(epsilon, 1-epsilon)\n",
    "    entropy = (-mask_list * torch.log(mask_list) - (1-mask_list) * torch.log(1-mask_list)).mean((1,2,3,4))\n",
    "    return entropy\n",
    "\n",
    "\n",
    "def get_mask_entropy(mask_list, threshold=0.01, epsilon=1e-5):\n",
    "    \"\"\"Obtain the pixel-wise entropy for each mask.\n",
    "    Args:\n",
    "        mask_list: each mask in mask_list should have the shape of [B, 1, H, W]\n",
    "        epsilon: to prevent NaN.\n",
    "\n",
    "    Returns:\n",
    "        entropy: shape [B,]\n",
    "    \"\"\"\n",
    "    def get_entropy(values):\n",
    "        return (-values * torch.log(values) - (1-values) * torch.log(1-values))\n",
    "    # Make sure that the format is correct:\n",
    "    assert isinstance(mask_list, list) or isinstance(mask_list, tuple)\n",
    "    shape = mask_list[0].shape\n",
    "    assert len(shape) == 4 and shape[1] == 1\n",
    "    \n",
    "    mask_list = torch.stack(mask_list, 1)\n",
    "    # Take mean over channel, height, and width\n",
    "    thresholded = torch.where(mask_list > threshold, 1, 0)\n",
    "    total_pixels = thresholded.sum((-3, -2, -1))\n",
    "    overall_values = torch.zeros(total_pixels.shape).to(total_pixels.device)\n",
    "    # Take mean over C, H, W\n",
    "    overall_values =  torch.where(total_pixels > 0, (thresholded * mask_list).sum((-3, -2, -1)) / total_pixels, overall_values)\n",
    "    entropy = get_entropy(overall_values.clamp(epsilon, 1-epsilon)).mean(1)\n",
    "    return entropy\n",
    "\n",
    "\n",
    "def get_neg_mask_overlap(\n",
    "    neg_mask: tuple,\n",
    "    mask_info=None,\n",
    "    is_penalize_lower=\"True\",\n",
    "    img=None,\n",
    "):\n",
    "    \"\"\"\n",
    "    Calculates penalty energy from negative masks overlapping with each other.\n",
    "\n",
    "    Args:\n",
    "        neg_mask: tuple of masks, each of shape [ensemble size, 1, H, W]\n",
    "        is_penalize_lower: if True or \"True\", will penalize that the sum is less than 1.\n",
    "            If \"False\" or False, will not. If \"obj:0.1\", will only penalize on the object locations (if n_channels==10), times coefficient of 0.1.\n",
    "        img: image, required if is_penalize_lower == \"obj\".\n",
    "\n",
    "    Returns:\n",
    "        tensor of shape [N,] representing each mask's energy\n",
    "    \"\"\"\n",
    "    # Stack all the masks on top of each other and sum over them\n",
    "    if mask_info is None or \"id_to_type\" not in mask_info:\n",
    "        neg_mask_obj = neg_mask\n",
    "    else:\n",
    "        neg_mask_obj = [neg_mask[i] for i in range(len(neg_mask)) if mask_info[\"id_to_type\"][i][1] == 1]\n",
    "    neg_mask_overlap = torch.cat(neg_mask_obj, 1).sum(1, keepdims=True)\n",
    "    if mask_info is not None and \"mask_exclude\" in mask_info:\n",
    "        neg_mask_overlap = neg_mask_overlap + mask_info[\"mask_exclude\"]\n",
    "    # Penalize the distance of the sum of masks compared to 1:\n",
    "    if is_penalize_lower == \"True\" or is_penalize_lower is True:\n",
    "        neg_mask_overlap = (neg_mask_overlap - 1).abs()\n",
    "    elif is_penalize_lower == \"False\" or is_penalize_lower is False:\n",
    "        neg_mask_overlap = (neg_mask_overlap - 1).clamp(0)\n",
    "    elif is_penalize_lower.startswith(\"obj\"):\n",
    "        penalize_lower_coef = eval(is_penalize_lower.split(\":\")[1]) if len(is_penalize_lower.split(\":\")) == 2 else 1\n",
    "        neg_mask_overlap_exceed = (neg_mask_overlap - 1).clamp(0)\n",
    "        assert len(img.shape) == 4\n",
    "        if img.shape[1] == 10:\n",
    "            obj_mask = (img[:, :1] != 1).float()\n",
    "            neg_mask_overlap_under = (obj_mask - neg_mask_overlap).clamp(0)\n",
    "            neg_mask_overlap = neg_mask_overlap_exceed + neg_mask_overlap_under * penalize_lower_coef\n",
    "        else:\n",
    "            neg_mask_overlap = neg_mask_overlap_exceed\n",
    "    else:\n",
    "        raise Exception(\"is_penalize_lower '{}' is not valid!\".format(is_penalize_lower))\n",
    "    return neg_mask_overlap.mean((1, 2, 3))\n",
    "\n",
    "\n",
    "def get_fine_neg_mask_overlap(neg_mask: tuple, threshold=0.01, epsilon=1e-5, mask_info=None, is_penalize_lower=True):\n",
    "    \"\"\"\n",
    "    Calculates penalty energy from negative masks overlapping with each other. Penalizes\n",
    "    overlap where sum of masks may not be greater than 1\n",
    "\n",
    "    Args:\n",
    "        neg_mask: tuple of masks, each of shape [ensemble size, 1, H, W]\n",
    "        is_penalize_lower: if True, will penalize that the sum is less than 1.\n",
    "\n",
    "    Returns:\n",
    "        tensor of shape [N,] representing each mask's energy\n",
    "    \"\"\"\n",
    "    # Stack all the masks on top of each other and sum over them\n",
    "    if mask_info is None:\n",
    "        neg_mask_obj = neg_mask\n",
    "    else:\n",
    "        neg_mask_obj = [neg_mask[i] for i in range(len(neg_mask)) if mask_info[\"id_to_type\"][i][1] == 1]\n",
    "    neg_mask_obj = torch.cat(neg_mask_obj, 1)\n",
    "    # Mask out pixels that are less than the threshold \n",
    "    neg_mask_thresh = (neg_mask_obj > threshold) * neg_mask_obj\n",
    "    neg_mask_overlap = neg_mask_thresh.sum(1, keepdims=True)\n",
    "    # Expand the sum of masks dimension\n",
    "    expanded_overlap = neg_mask_overlap.expand(-1, neg_mask_thresh.shape[1], -1, -1)\n",
    "    other_mask = expanded_overlap - neg_mask_thresh\n",
    "    \n",
    "    overlapping = torch.where((other_mask > threshold).logical_and(neg_mask_thresh > threshold), 1.0, 0.0)\n",
    "    mean_mask = neg_mask_thresh.mean((2, 3)) # Take mean over H, W\n",
    "    mean_overlap = (neg_mask_thresh * overlapping).mean((2, 3))\n",
    "    normalized_overlap = mean_overlap / mean_mask.clamp(min=epsilon)\n",
    "    assert torch.all(normalized_overlap <= 1.0)\n",
    "    mean_overlap = (normalized_overlap).mean(1) # Take mean over EBM's\n",
    "    return mean_overlap\n",
    "\n",
    "\n",
    "def get_neg_mask_exceed(neg_mask: tuple, pos_all_mask: Union[torch.Tensor, tuple]):\n",
    "    \"\"\"\n",
    "    Calculates penalty energy from a negative mask appearing where there is no\n",
    "    object in reality (aka, whenever negative mask isn't a subset of ground\n",
    "    truth mask).\n",
    "\n",
    "    :param neg_mask: tuple of masks, each of shape [N, 1, H, W]\n",
    "    :param pos_all_mask: ground truth mask(s), where each [H, W] contains the\n",
    "    mask for *all* objects in the image. Either tensor of shape [N, 1, H, W] or\n",
    "    tuple. If tuple, it doesn't represent all the objects?\n",
    "    :returns: tensor of shape [N,] representing each mask's energy\n",
    "    \"\"\"\n",
    "\n",
    "    # TODO pos_all_mask tuple\n",
    "    neg_mask_stack = torch.cat(neg_mask, 1).sum(1, keepdims=True).clamp(max=1)\n",
    "\n",
    "    exceeding_regions = neg_mask_stack * pos_all_mask.logical_not()\n",
    "\n",
    "    return exceeding_regions.mean((1, 2, 3))\n",
    "\n",
    "\n",
    "def get_neg_mask_exceed_energy(neg_mask, pos_all_mask):\n",
    "    \"\"\"\n",
    "    Calculates \"penalty\" energy given a model-predicted negative mask\n",
    "    and a positive ground-truth mask, which increases when:\n",
    "\n",
    "    1. Two negative masks overlap\n",
    "    2. A negative mask appears in a place where there is no object in reality.\n",
    "    \"\"\"\n",
    "\n",
    "    if isinstance(pos_all_mask, tuple):\n",
    "        assert len(neg_mask) == 2\n",
    "        neg_mask_exceed = torch.cat([torch.clamp(neg_mask[0] - pos_all_mask[0], min=0),\n",
    "                                     torch.clamp(neg_mask[1] - pos_all_mask[1], min=0)], 1).sum(1, keepdims=True)\n",
    "    else:\n",
    "        neg_mask_sum = torch.cat(neg_mask, 1).sum(1, keepdims=True)\n",
    "        # Calculate number of times mask exceeds boundary of positive ground truth\n",
    "        neg_mask_exceed = torch.clamp(neg_mask_sum - pos_all_mask, min=0)\n",
    "    neg_mask_exceed_energy = neg_mask_exceed.mean((1,2,3))\n",
    "    return neg_mask_exceed_energy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. GNN:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 GNLayer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GNLayer(MessagePassing):\n",
    "    \"\"\"One GN layer of edge update followed by node update.\"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        edge_attr_size,\n",
    "        n_neurons,\n",
    "        mlp_n_layers,\n",
    "        aggr_mode=\"max\",\n",
    "        output_size=None,\n",
    "        activation=\"relu\",\n",
    "        is_output_edge_feature=False,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Initialization for the GNLayer.\n",
    "        Args:\n",
    "            input_size: number of features for each node.\n",
    "            edge_attr_size: number of features for edge attribute.\n",
    "            n_neurons: number of neurons for each layer of the node_mlp and edge_mlp.\n",
    "            mlp_n_layers: number of layers for the node_mlp and edge_mlp.\n",
    "            output_size: output size. If None, will use edge_attr_size.\n",
    "            activation: activation for the hidden layers for the node_mlp and edge_mlp.\n",
    "        \"\"\"\n",
    "        super(GNLayer, self).__init__(aggr=aggr_mode)\n",
    "        self.is_output_edge_feature = is_output_edge_feature\n",
    "        self.node_mlp = MLP(input_size=n_neurons + input_size,\n",
    "                            n_neurons=n_neurons,\n",
    "                            n_layers=mlp_n_layers,\n",
    "                            output_size=edge_attr_size if output_size is None else output_size,\n",
    "                            activation=activation,\n",
    "                           )\n",
    "        self.edge_mlp = MLP(input_size=2 * input_size + edge_attr_size,\n",
    "                            n_neurons=n_neurons,\n",
    "                            n_layers=mlp_n_layers,\n",
    "                            activation=activation,\n",
    "                           )\n",
    "\n",
    "    def forward(self, x, edge_index, edge_attr):\n",
    "        \"\"\"Main forward function\"\"\"\n",
    "        # x has shape [N, in_channels]\n",
    "        # edge_index has shape [2, E]\n",
    "\n",
    "        # Computes the aggregated results after summing over messages from neighboring edges for each node:\n",
    "        out = self.propagate(edge_index, size=(x.size(0), x.size(0)), x=x, edge_attr=edge_attr)\n",
    "        # Use an MLP on the concatenation of original node features and the aggregated messages:\n",
    "        out = self.node_mlp(torch.cat([x, out], 1))\n",
    "        if self.is_output_edge_feature:\n",
    "            return out, self.edge_feature\n",
    "        else:\n",
    "            return out\n",
    "\n",
    "    def message(self, x_i, x_j, edge_attr):\n",
    "        \"\"\"Computes the message given features of the source node (x_i) and features of the target node.\"\"\"\n",
    "        # x_i has shape [E, in_channels]\n",
    "        # x_j has shape [E, in_channels]\n",
    "        # edge_attr has shape [E, in_channels]\n",
    "        out = torch.cat([x_i, x_j, edge_attr], dim=1)  # [E, 2 * in_channels + edge_attr_size]\n",
    "        out = self.edge_mlp(out)\n",
    "        if self.is_output_edge_feature:\n",
    "            self.edge_feature = out\n",
    "        return out\n",
    "\n",
    "    def update(self, aggr_out):\n",
    "        \"\"\"The update function, after summing over the messages from the neighboring edges.\"\"\"\n",
    "        # aggr_out has shape [N, out_channels]\n",
    "        return aggr_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 GNNs:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2.1 GNN2:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GNN2(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        edge_attr_size,\n",
    "        n_GN_layers,\n",
    "        n_neurons=32,\n",
    "        GNN_output_size=8,\n",
    "        mlp_n_layers=2,\n",
    "        normalize=\"None\",\n",
    "        activation=\"relu\",\n",
    "        recurrent=False,\n",
    "        aggr_mode=\"max\",\n",
    "        is_output_edge_feature=False,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            input_size: number of features for each node.\n",
    "            edge_attr_size: number of features for edge attribute.\n",
    "            n_GN_layers: number of GNLayers.\n",
    "            n_neurons: number of neurons for each layer of the node_mlp and edge_mlp.\n",
    "            GNN_output_size: output_size of each GNLayer. If None, will use edge_attr_size.\n",
    "            mlp_n_layers: number of layers for the node_mlp and edge_mlp.\n",
    "            normalize: normalization mode. Choose from \"None\", \"layer\", \"batch\".\n",
    "            activation: activation for the hidden layers for the node_mlp and edge_mlp.\n",
    "        \"\"\"\n",
    "        super(GNN2, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.n_GN_layers = n_GN_layers\n",
    "        self.GNN_output_size = GNN_output_size\n",
    "        self.normalize = normalize\n",
    "        self.aggr_mode = aggr_mode\n",
    "        self.is_output_edge_feature = is_output_edge_feature\n",
    "        for i in range(1, self.n_GN_layers + 1):\n",
    "            setattr(self, \"layer_{}\".format(i),\n",
    "                    GNLayer(\n",
    "                    input_size=input_size if i == 1 else GNN_output_size,\n",
    "                    edge_attr_size=edge_attr_size,\n",
    "                    aggr_mode=aggr_mode,\n",
    "                    n_neurons=n_neurons,\n",
    "                    mlp_n_layers=mlp_n_layers,\n",
    "                    output_size=GNN_output_size,\n",
    "                    activation=activation,\n",
    "                    is_output_edge_feature=is_output_edge_feature if i == self.n_GN_layers else False,\n",
    "            ))\n",
    "            if i != self.n_GN_layers:\n",
    "                if self.normalize == \"layer\":\n",
    "                    setattr(self, \"norm_{}\".format(i), nn.LayerNorm(GNN_output_size))\n",
    "                elif self.normalize.startswith(\"gn\"):\n",
    "                    n_groups = eval(self.normalize.split(\"-\")[1])\n",
    "                    setattr(self, \"norm_{}\".format(i), nn.GroupNorm(n_groups, GNN_output_size, affine=True))\n",
    "                elif self.normalize == \"batch\":\n",
    "                    setattr(self, \"norm_{}\".format(i), nn.BatchNorm1d(GNN_output_size))\n",
    "                elif self.normalize == \"None\":\n",
    "                    pass\n",
    "                else:\n",
    "                    raise Exception(\"Normalize '{}' is not valid!\".format(self.normalize))\n",
    "        self.rnn = None \n",
    "        if recurrent:\n",
    "            # num input features, num hidden features. Always uses bias by default.\n",
    "            # Batch size is the number of nodes in the graph\n",
    "            self.rnn = nn.GRUCell(GNN_output_size, GNN_output_size)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        data,\n",
    "        rnn_hx=None,\n",
    "    ):\n",
    "        x = data.x\n",
    "        for i in range(1, self.n_GN_layers + 1):\n",
    "            if i == self.n_GN_layers and self.is_output_edge_feature:\n",
    "                x, edge_features = getattr(self, \"layer_{}\".format(i))(x, edge_index=data.edge_index, edge_attr=data.edge_attr)\n",
    "            else:\n",
    "                x = getattr(self, \"layer_{}\".format(i))(x, edge_index=data.edge_index, edge_attr=data.edge_attr)\n",
    "            if i != self.n_GN_layers and (self.normalize in [\"layer\", \"batch\"] or self.normalize.startswith(\"gn\")):\n",
    "                x = getattr(self, \"norm_{}\".format(i))(x)\n",
    "\n",
    "        embed_hx_new = None\n",
    "        if self.rnn is not None:\n",
    "            # Pass in input, hidden state\n",
    "            x = self.rnn(x, rnn_hx)\n",
    "            embed_hx_new = x.clone()\n",
    "        if self.is_output_edge_feature:\n",
    "            return (x, edge_features), embed_hx_new\n",
    "        else:\n",
    "            return x, embed_hx_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2.2 GNN_energy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GNN_energy(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        mode,\n",
    "        is_zgnn_node,\n",
    "        edge_attr_size,\n",
    "        aggr_mode,\n",
    "        n_GN_layers,\n",
    "        n_neurons=32,\n",
    "        GNN_output_size=8,\n",
    "        mlp_n_layers=2,\n",
    "        gnn_normalization_type=\"None\",\n",
    "        activation=\"relu\",\n",
    "        recurrent=False,\n",
    "        cnn_output_size=8,\n",
    "        cnn_is_spec_norm=True,\n",
    "        cnn_normalization_type=\"None\",\n",
    "        cnn_channel_base=64,\n",
    "        cnn_aggr_mode=\"sum\",\n",
    "        c_repr_dim=8,\n",
    "        z_dim=8,\n",
    "        zgnn_dim=1,\n",
    "        distance_loss_type=\"Jaccard\",\n",
    "        pooling_type=\"gated\",\n",
    "        pooling_dim=32,\n",
    "        is_x=False,\n",
    "    ):\n",
    "        \"\"\"Combine multiple modalities (w, z, c, zgnn, wtarget) together via CNN and GNN, \n",
    "        and return a single energy.\n",
    "        \n",
    "        Args:\n",
    "            mode: choose from \"concat\", \"softmax\"\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.mode = mode\n",
    "        self.is_zgnn_node = is_zgnn_node\n",
    "        self.is_x = is_x\n",
    "        if self.mode == \"softmax\":\n",
    "            assert zgnn_dim == 1\n",
    "            self.softmax_coef = nn.Parameter(torch.ones(1) * 0.2)\n",
    "            self.in_channels = 1\n",
    "        else:\n",
    "            self.in_channels = 1\n",
    "        self.gnn_input_size = cnn_output_size + z_dim + c_repr_dim\n",
    "        if self.is_zgnn_node:\n",
    "            self.gnn_input_size += zgnn_dim\n",
    "            self.in_channels += 1\n",
    "        else:\n",
    "            if self.is_x:\n",
    "                self.in_channels = 10\n",
    "        self.cnn_aggr_mode = cnn_aggr_mode\n",
    "        self.edge_attr_size = edge_attr_size\n",
    "        self.aggr_mode = aggr_mode\n",
    "        self.n_GN_layers = n_GN_layers\n",
    "        self.n_neurons = n_neurons\n",
    "        self.GNN_output_size = GNN_output_size\n",
    "        self.mlp_n_layers = mlp_n_layers\n",
    "        self.gnn_normalization_type = gnn_normalization_type\n",
    "        self.activation = activation\n",
    "        self.recurrent = recurrent\n",
    "        self.cnn_output_size = cnn_output_size\n",
    "        self.cnn_is_spec_norm = cnn_is_spec_norm\n",
    "        self.cnn_normalization_type = cnn_normalization_type\n",
    "        self.cnn_channel_base = cnn_channel_base\n",
    "        self.cnn_aggr_mode = cnn_aggr_mode\n",
    "        self.c_repr_dim = c_repr_dim\n",
    "        self.z_dim = z_dim\n",
    "        self.zgnn_dim = zgnn_dim\n",
    "        self.distance_loss_type = distance_loss_type\n",
    "        self.pooling_type = pooling_type\n",
    "        self.pooling_dim = pooling_dim\n",
    "\n",
    "        if distance_loss_type == \"mse\":\n",
    "            self.loss_fun = lambda x, y: nn.MSELoss(reduction=\"none\")(x, y).mean((-3,-2,-1))\n",
    "        elif distance_loss_type == \"Jaccard\":\n",
    "            self.loss_fun = lambda x, y: get_soft_Jaccard_distance(x, y, dim=(-3,-2,-1))\n",
    "        else:\n",
    "            raise\n",
    "        self.gnn = GNN2(\n",
    "            input_size=self.gnn_input_size,\n",
    "            edge_attr_size=edge_attr_size,\n",
    "            aggr_mode=aggr_mode,\n",
    "            n_GN_layers=n_GN_layers,\n",
    "            n_neurons=n_neurons,\n",
    "            GNN_output_size=GNN_output_size,\n",
    "            mlp_n_layers=mlp_n_layers,\n",
    "            normalize=gnn_normalization_type,\n",
    "            activation=activation,\n",
    "            recurrent=recurrent,\n",
    "            is_output_edge_feature=True if pooling_type == \"gated\" else False,\n",
    "        )\n",
    "        self.cnn = nn.Sequential(\n",
    "            CResBlock(self.in_channels, cnn_channel_base, downsample=True, is_spec_norm=cnn_is_spec_norm, c_repr_mode=\"None\", z_mode=\"None\", act_name=activation, normalization_type=cnn_normalization_type),\n",
    "            CResBlock(cnn_channel_base, cnn_channel_base, is_spec_norm=cnn_is_spec_norm, c_repr_mode=\"None\", z_mode=\"None\", act_name=activation, normalization_type=cnn_normalization_type),\n",
    "            CResBlock(cnn_channel_base, cnn_channel_base*2, downsample=True, is_spec_norm=cnn_is_spec_norm, c_repr_mode=\"None\", z_mode=\"None\", act_name=activation, normalization_type=cnn_normalization_type),\n",
    "            CResBlock(cnn_channel_base*2, cnn_channel_base*2, is_spec_norm=cnn_is_spec_norm, c_repr_mode=\"None\", z_mode=\"None\", act_name=activation, normalization_type=cnn_normalization_type),\n",
    "            CResBlock(cnn_channel_base*2, cnn_channel_base*2, downsample=True, is_spec_norm=cnn_is_spec_norm, c_repr_mode=\"None\", z_mode=\"None\", act_name=activation, normalization_type=cnn_normalization_type),\n",
    "            CResBlock(cnn_channel_base*2, cnn_channel_base*2, is_spec_norm=cnn_is_spec_norm, c_repr_mode=\"None\", z_mode=\"None\", act_name=activation, normalization_type=cnn_normalization_type),\n",
    "        )\n",
    "        self.cnn_mlp = MLP(\n",
    "            cnn_channel_base*2,\n",
    "            n_neurons=cnn_channel_base,\n",
    "            n_layers=1,\n",
    "            activation=activation,\n",
    "            output_size=cnn_output_size,\n",
    "        )\n",
    "        if self.z_dim > 0:\n",
    "            self.g_z = MLP(\n",
    "                z_dim,\n",
    "                n_neurons=z_dim,\n",
    "                n_layers=1,\n",
    "                activation=activation,\n",
    "                output_size=z_dim,\n",
    "            )\n",
    "        if self.pooling_type == \"mean\":\n",
    "            self.gnn_mlp = MLP(\n",
    "                GNN_output_size,\n",
    "                n_neurons=GNN_output_size,\n",
    "                n_layers=1,\n",
    "                activation=\"relu\",\n",
    "                output_size=1,\n",
    "            )\n",
    "        elif self.pooling_type == \"gated\":\n",
    "            self.node_attn = nn.Linear(GNN_output_size, 1)\n",
    "            self.node_pooling = nn.Linear(GNN_output_size, pooling_dim)\n",
    "            self.edge_attn = nn.Linear(n_neurons, 1)\n",
    "            self.edge_pooling = nn.Linear(n_neurons, pooling_dim)\n",
    "            self.gnn_mlp = MLP(\n",
    "                pooling_dim * 2,\n",
    "                n_neurons=pooling_dim,\n",
    "                n_layers=2,\n",
    "                activation=\"relu\",\n",
    "                last_layer_linear=False,\n",
    "                output_size=1,\n",
    "            )\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "\n",
    "    def cnn_forward(self, input):\n",
    "        out = self.cnn(input)\n",
    "        if self.cnn_aggr_mode == \"sum\":\n",
    "            out = out.view(out.shape[0], out.shape[1], -1).sum(2)\n",
    "        elif self.cnn_aggr_mode == \"max\":\n",
    "            out = out.view(out.shape[0], out.shape[1], -1).max(2)[0]\n",
    "        elif self.cnn_aggr_mode == \"mean\":\n",
    "            out = out.view(out.shape[0], out.shape[1], -1).mean(2)\n",
    "        else:\n",
    "            raise\n",
    "        out = self.cnn_mlp(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "    def get_pyg_data(self, w, z, c, zgnn, wtarget, batch_shape, x=None):\n",
    "        def get_edge_index_aug(n_ebms, batch_size):\n",
    "            edge_matrix = np.stack(np.meshgrid(np.arange(n_ebms), np.arange(n_ebms)), -1).reshape(-1,2)\n",
    "            mask = edge_matrix[:,0] != edge_matrix[:,1]\n",
    "            edge_index = torch.LongTensor(edge_matrix[mask].T)\n",
    "            edge_index_aug = edge_index[:,None]\n",
    "            edge_index_aug = edge_index_aug + (torch.arange(batch_size) * n_ebms)[None,:,None]\n",
    "            edge_index_aug = edge_index_aug.reshape(2, -1)\n",
    "            return edge_index_aug\n",
    "\n",
    "        n_ebms = len(w)\n",
    "        batch_size_all = np.prod(batch_shape)\n",
    "        device = w[0].device\n",
    "        if self.mode == \"concat\":\n",
    "            w_embed = torch.stack([self.cnn_forward(torch.cat([w_ele, wtarget], 1) if self.is_zgnn_node else w_ele * x if self.is_x and x is not None else w_ele) for w_ele in w], 1)  # each has shape [B_task * B_example, n_ebms, cnn_output_size]\n",
    "        elif self.mode == \"softmax\":\n",
    "            w_embed = torch.stack([self.cnn_forward(w_ele) for w_ele in w], 1)  # each has shape [B_task * B_example, n_ebms, cnn_output_size]\n",
    "        else:\n",
    "            raise Exception(\"mode '{}' is not valid!\".format(self.mode))\n",
    "        c_embed = torch.stack(c, 1)  # [1, n_ebms, c_repr_dim]\n",
    "        c_embed = c_embed.expand(w_embed.shape[0], n_ebms, c_embed.shape[-1])     # [B_task * B_example, n_ebms, c_repr_dim]\n",
    "\n",
    "        if self.z_dim > 0:\n",
    "            z_embed = self.g_z(torch.stack(z, 1))  # [B_task * B_example, n_ebms, z_dim]\n",
    "            all_embed = [w_embed, z_embed, c_embed]\n",
    "        else:\n",
    "            all_embed = [w_embed, c_embed]\n",
    "        if self.is_zgnn_node:\n",
    "            zgnn_node = zgnn[0][:,None].expand(*batch_shape, *zgnn[0].shape[-2:])  # [B_task, B_example, n_ebms, Z_node_size]\n",
    "            zgnn_node_embed = zgnn_node.reshape(-1, *zgnn_node.shape[-2:])         # [B_task * B_example, n_ebms, Z_node_size]\n",
    "            all_embed.append(zgnn_node_embed)\n",
    "        data_x = torch.cat(all_embed, -1)   # [B_task * B_example, n_ebms, gnn_input_size]\n",
    "        data_edge_index = get_edge_index_aug(n_ebms, batch_size_all).to(device)    # [2, B_task * B_example * n_edges], n_edges = n_ebm * (n_ebm-1)\n",
    "        zgnn_edge_embed = zgnn[1][:,None].expand(*batch_shape, *zgnn[1].shape[-2:])# zgnn[1]: [B_task, n_edges, edge_attr_size], zgnn_edge_embed: [B_task, B_example, n_edges, edge_attr_size]\n",
    "        data_edge_attr = zgnn_edge_embed.reshape(-1, zgnn_edge_embed.shape[-1])    # [B_task * B_example * n_edges, edge_attr_size]\n",
    "\n",
    "        data = Data(x=data_x.view(-1, data_x.shape[-1]), edge_index=data_edge_index, edge_attr=data_edge_attr)\n",
    "        return data\n",
    "\n",
    "\n",
    "    def forward(self, w, z, c, zgnn, wtarget, batch_shape, x=None):\n",
    "        data = self.get_pyg_data(w, z, c, zgnn, wtarget, batch_shape=batch_shape, x=x)\n",
    "        assert data.x.shape[-1] == self.gnn.input_size\n",
    "\n",
    "        if self.pooling_type == \"mean\":\n",
    "            gnn_out, _ = self.gnn(data)\n",
    "            gnn_out = self.gnn_mlp(gnn_out).view(*batch_shape, -1)\n",
    "            E_all = gnn_out.mean(-1, keepdims=True)  # [B_task, B_example, 1]\n",
    "        elif self.pooling_type == \"gated\":\n",
    "            n_nodes = len(w)\n",
    "            n_edges = n_nodes * (n_nodes-1)\n",
    "            (node_feature_out, edge_feature_out), _ = self.gnn(data)\n",
    "            node_feature_out = node_feature_out.view(*batch_shape, n_nodes, -1)  # [B_task, B_example, n_nodes, GNN_output_size]\n",
    "            edge_feature_out = edge_feature_out.view(*batch_shape, n_edges, -1)  # [B_task, B_example, n_nodes, n_neurons]\n",
    "\n",
    "            node_pool = self.node_pooling(node_feature_out * self.node_attn(node_feature_out)).sum(-2) # [B_task, B_example, pool_dim]\n",
    "            edge_pool = self.edge_pooling(edge_feature_out * self.edge_attn(edge_feature_out)).sum(-2) # [B_task, B_example, pool_dim]\n",
    "            pool = torch.cat([node_pool, edge_pool], -1)\n",
    "            E_all = self.gnn_mlp(pool)\n",
    "        if self.mode == \"softmax\":\n",
    "            prob = zgnn[0] / (zgnn[0].abs().sum(1, keepdims=True) + 1e-5)  # [B_task, n_ebms, 1]\n",
    "            entropy = - (prob * torch.log(prob.clamp(1e-5))).sum(1, keepdims=True)\n",
    "            w_reshape = torch.stack([w_ele.view(*batch_shape, *w_ele.shape[-3:]) for w_ele in w], 1)  # [B_task, n_ebms, B_example, 1, H, W]\n",
    "            w_weighted = (w_reshape * prob[...,None,None,None]).sum(1)  # \n",
    "            w_weighted = w_weighted.view(-1, *w_weighted.shape[-3:])\n",
    "            E_fit = self.loss_fun(w_weighted, wtarget).view(*batch_shape, 1) * self.softmax_coef\n",
    "            entropy_expand = entropy.expand(E_fit.shape) * self.softmax_coef / 3\n",
    "            E_all = E_all + E_fit\n",
    "        return E_all\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"GNN_energy\"}\n",
    "        model_dict[\"mode\"] = self.mode\n",
    "        model_dict[\"is_zgnn_node\"] = self.is_zgnn_node\n",
    "        model_dict[\"edge_attr_size\"] = self.edge_attr_size\n",
    "        model_dict[\"aggr_mode\"] = self.aggr_mode\n",
    "        model_dict[\"n_GN_layers\"] = self.n_GN_layers\n",
    "        model_dict[\"n_neurons\"] = self.n_neurons\n",
    "        model_dict[\"GNN_output_size\"] = self.GNN_output_size\n",
    "        model_dict[\"mlp_n_layers\"] = self.mlp_n_layers\n",
    "        model_dict[\"gnn_normalization_type\"] = self.gnn_normalization_type\n",
    "        model_dict[\"activation\"] = self.activation\n",
    "        model_dict[\"recurrent\"] = self.recurrent\n",
    "        model_dict[\"cnn_output_size\"] = self.cnn_output_size\n",
    "        model_dict[\"cnn_is_spec_norm\"] = self.cnn_is_spec_norm\n",
    "        model_dict[\"cnn_normalization_type\"] = self.cnn_normalization_type\n",
    "        model_dict[\"cnn_channel_base\"] = self.cnn_channel_base\n",
    "        model_dict[\"cnn_aggr_mode\"] = self.cnn_aggr_mode\n",
    "        model_dict[\"c_repr_dim\"] = self.c_repr_dim\n",
    "        model_dict[\"z_dim\"] = self.z_dim\n",
    "        model_dict[\"zgnn_dim\"] = self.zgnn_dim\n",
    "        model_dict[\"distance_loss_type\"] = self.distance_loss_type\n",
    "        model_dict[\"pooling_type\"] = self.pooling_type\n",
    "        model_dict[\"pooling_dim\"] = self.pooling_dim\n",
    "        model_dict[\"is_x\"] = self.is_x\n",
    "        model_dict[\"state_dict\"] = to_cpu(self.state_dict())\n",
    "        return model_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    device = \"cuda:0\"\n",
    "    batch_shape = (20,6)\n",
    "    w = (torch.rand(120,1,8,8), torch.rand(120,1,8,8), torch.rand(120,1,8,8), torch.rand(120,1,8,8))\n",
    "    z = (torch.rand(120,5),torch.rand(120,5),torch.rand(120,5),torch.rand(120,5))\n",
    "    c = (torch.rand(1,8), torch.rand(1,8), torch.rand(1,8), torch.rand(1,8))\n",
    "    zgnn = (torch.rand(20,4,1), torch.rand(20,12,8))\n",
    "    wtarget = None #torch.rand(12,1,8,8)\n",
    "    w, z, c, zgnn, wtarget = to_device_recur((w, z, c, zgnn, wtarget), device=device)\n",
    "\n",
    "    self = GNN_energy(\n",
    "        mode=\"concat\",\n",
    "        is_zgnn_node=False,\n",
    "        aggr_mode=\"add\",\n",
    "        edge_attr_size=8,\n",
    "        n_GN_layers=2,\n",
    "        z_dim=5,\n",
    "        zgnn_dim=1,\n",
    "        distance_loss_type=\"Jaccard\",\n",
    "        pooling_type=\"gated\",\n",
    "        pooling_dim=32,\n",
    "        cnn_is_spec_norm=True,\n",
    "    ).to(device)\n",
    "    out = self(w, z, c, zgnn, wtarget, batch_shape)\n",
    "    self.model_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. LambdaNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.1 LambdaLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MIT License\n",
    "# Adapted from https://github.com/lucidrains/lambda-networks\n",
    "\n",
    "from torch import einsum\n",
    "from einops import rearrange\n",
    "\n",
    "# helpers functions\n",
    "\n",
    "def exists(val):\n",
    "    return val is not None\n",
    "\n",
    "def default(val, d):\n",
    "    return val if exists(val) else d\n",
    "\n",
    "def calc_rel_pos(n):\n",
    "    pos = torch.meshgrid(torch.arange(n), torch.arange(n))\n",
    "    pos = rearrange(torch.stack(pos), 'n i j -> (i j) n')  # [n*n, 2] pos[n] = (i, j)\n",
    "    rel_pos = pos[None, :] - pos[:, None]                  # [n*n, n*n, 2] rel_pos[n, m] = (rel_i, rel_j)\n",
    "    rel_pos += n - 1                                       # shift value range from [-n+1, n-1] to [0, 2n-2]\n",
    "    return rel_pos\n",
    "\n",
    "# lambda layer\n",
    "\n",
    "class LambdaLayer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        # Dimension of the query\n",
    "        dim_q,\n",
    "        # Dimension of the context\n",
    "        dim_c,\n",
    "        *,\n",
    "        dim_k,\n",
    "        # The length of the context\n",
    "        m = None,\n",
    "        r = None,\n",
    "        heads = 4,\n",
    "        # Dimension of the value if heads=1\n",
    "        dim_out = None,\n",
    "        dim_u = 1,\n",
    "        use_relpos=True):\n",
    "        super().__init__()\n",
    "        dim_out = default(dim_out, dim_c)\n",
    "        self.u = dim_u # intra-depth dimension\n",
    "        self.heads = heads\n",
    "\n",
    "        assert (dim_out % heads) == 0, 'values dimension must be divisible by number of heads for multi-head query'\n",
    "        dim_v = dim_out // heads\n",
    "        self.to_q = nn.Conv2d(dim_q, dim_k * heads, 1, bias = False)\n",
    "        self.to_k = nn.Conv2d(dim_c, dim_k * dim_u, 1, bias = False)\n",
    "        self.to_v = nn.Conv2d(dim_c, dim_v * dim_u, 1, bias = False)\n",
    "\n",
    "        self.norm_q = nn.BatchNorm2d(dim_k * heads)\n",
    "        self.norm_v = nn.BatchNorm2d(dim_v * dim_u)\n",
    "\n",
    "        self.local_contexts = exists(r)\n",
    "        self.use_relpos = use_relpos\n",
    "        if exists(r):\n",
    "            assert (r % 2) == 1, 'Receptive kernel size should be odd'\n",
    "            self.pos_conv = nn.Conv3d(dim_u, dim_k, (1, r, r), padding = (0, r // 2, r // 2))\n",
    "        elif self.use_relpos:\n",
    "            assert exists(m), 'You must specify the window size (m=h=w)'\n",
    "            rel_lengths = 2 * m - 1\n",
    "            self.rel_pos_emb = nn.Parameter(torch.randn(rel_lengths, rel_lengths, dim_k, dim_u))\n",
    "            self.rel_pos = calc_rel_pos(m)\n",
    "\n",
    "    def forward(self, x, context):\n",
    "        \"\"\"Take in both x and the context.\"\"\"\n",
    "        b, c, hh, ww, u, h = *context.shape, self.u, self.heads\n",
    "\n",
    "        q = self.to_q(x)\n",
    "        k = self.to_k(context)\n",
    "        v = self.to_v(context)\n",
    "\n",
    "        q = self.norm_q(q)\n",
    "        v = self.norm_v(v)\n",
    "\n",
    "        q = rearrange(q, 'b (h k) hh ww -> b h k (hh ww)', h = h)\n",
    "        k = rearrange(k, 'b (u k) hh ww -> b u k (hh ww)', u = u)\n",
    "        v = rearrange(v, 'b (u v) hh ww -> b u v (hh ww)', u = u)\n",
    "\n",
    "        k = k.softmax(dim=-1)\n",
    "\n",
    "        λc = einsum('b u k m, b u v m -> b k v', k, v)\n",
    "        Yc = einsum('b h k n, b k v -> b h v n', q, λc)\n",
    "\n",
    "        if self.local_contexts:\n",
    "            v = rearrange(v, 'b u v (hh ww) -> b u v hh ww', hh = hh, ww = ww)\n",
    "            λp = self.pos_conv(v)\n",
    "            Yp = einsum('b h k n, b k v n -> b h v n', q, λp.flatten(3))\n",
    "        elif self.use_relpos:\n",
    "            n, m = self.rel_pos.unbind(dim = -1)\n",
    "            rel_pos_emb = self.rel_pos_emb[n, m]\n",
    "            λp = einsum('n m k u, b u v m -> b n k v', rel_pos_emb, v)\n",
    "            Yp = einsum('b h k n, b n k v -> b h v n', q, λp)\n",
    "        else:\n",
    "            # Use the shape of the queries\n",
    "            _, _, hh, ww = x.shape\n",
    "            Yp = torch.zeros(Yc.shape, device=Yc.device)\n",
    "\n",
    "        Y = Yc + Yp\n",
    "        out = rearrange(Y, 'b h v (hh ww) -> b (h v) hh ww', hh = hh, ww = ww)\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2 Test LambdaLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    # Batch size, num_channels, height, width\n",
    "    x = torch.randn(1, 32, 64, 64)\n",
    "    c = torch.clone(x)\n",
    "    layer = LambdaLayer(\n",
    "        dim_q = 32,\n",
    "        dim_c = 32,       # channels going in\n",
    "        dim_out = 32,   # channels out\n",
    "        m = 64,         # size of the receptive window - max(height, width)\n",
    "        dim_k = 16,     # key dimension\n",
    "        heads = 4,      # number of heads, for multi-query\n",
    "        dim_u = 1       # 'intra-depth' dimension\n",
    "    )\n",
    "    out = layer(x, c)\n",
    "    print(out.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    # Test query having different dimension from context, as well as different\n",
    "    # height and width\n",
    "    x = torch.randn(1, 10, 1, 50)\n",
    "    c = torch.randn(1, 32, 64, 64)\n",
    "    layer = LambdaLayer(\n",
    "        dim_q = 10,\n",
    "        dim_c = 32,       # channels going in\n",
    "        dim_out = 32,   # channels out\n",
    "        m = 64,         # size of the receptive window - max(height, width)\n",
    "        dim_k = 16,     # key dimension\n",
    "        heads = 4,      # number of heads, for multi-query\n",
    "        dim_u = 1,       # 'intra-depth' dimension\n",
    "        use_relpos=False # Using relpos embedding when query isn't same shape as \n",
    "                        # context isn't implemented\n",
    "    )\n",
    "    out = layer(x, c)\n",
    "    print(out.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. MONet:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.1 AttentionNet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# License: MIT\n",
    "# Adapted from https://github.com/stelzner/monet\n",
    "\n",
    "def double_conv(in_channels, out_channels, activation):\n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(in_channels, out_channels, 3, padding=1),\n",
    "        nn.BatchNorm2d(out_channels),\n",
    "        get_activation(activation, inplace=True),\n",
    "        nn.Conv2d(out_channels, out_channels, 3, padding=1),\n",
    "        nn.BatchNorm2d(out_channels),\n",
    "        get_activation(activation, inplace=True),\n",
    "    )\n",
    "\n",
    "\n",
    "class UNet(nn.Module):\n",
    "    def __init__(self, n_blocks, in_channels, out_channels, channel_base=64, activation=\"relu\"):\n",
    "        super().__init__()\n",
    "        self.n_blocks = n_blocks\n",
    "        self.down_convs = nn.ModuleList()\n",
    "        cur_in_channels = in_channels\n",
    "        for i in range(n_blocks):\n",
    "            self.down_convs.append(double_conv(cur_in_channels,\n",
    "                                               channel_base * 2**i,\n",
    "                                               activation=activation,\n",
    "                                              ))\n",
    "            cur_in_channels = channel_base * 2**i\n",
    "\n",
    "        self.tconvs = nn.ModuleList()\n",
    "        for i in range(n_blocks-1, 0, -1):\n",
    "            self.tconvs.append(nn.ConvTranspose2d(channel_base * 2**i,\n",
    "                                                  channel_base * 2**(i-1),\n",
    "                                                  2, stride=2))\n",
    "\n",
    "        self.up_convs = nn.ModuleList()\n",
    "        for i in range(n_blocks-2, -1, -1):\n",
    "            self.up_convs.append(double_conv(channel_base * 2**(i+1),\n",
    "                                             channel_base * 2**i,\n",
    "                                             activation=activation,\n",
    "                                            ))\n",
    "\n",
    "        self.final_conv = nn.Conv2d(channel_base, out_channels, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        intermediates = []\n",
    "        cur = x\n",
    "        for down_conv in self.down_convs[:-1]:\n",
    "            cur = down_conv(cur)\n",
    "            intermediates.append(cur)\n",
    "            cur = nn.MaxPool2d(2)(cur)\n",
    "\n",
    "        cur = self.down_convs[-1](cur)\n",
    "\n",
    "        for i in range(self.n_blocks-1):\n",
    "            cur = self.tconvs[i](cur)\n",
    "            cur = torch.cat((cur, intermediates[-i -1]), 1)\n",
    "            cur = self.up_convs[i](cur)\n",
    "\n",
    "        return self.final_conv(cur)\n",
    "\n",
    "\n",
    "class AttentionNet(nn.Module):\n",
    "    def __init__(self, in_channels, n_blocks, channel_base, activation=\"relu\"):\n",
    "        super().__init__()\n",
    "        self.unet = UNet(n_blocks=n_blocks,\n",
    "                         in_channels=in_channels+1, #->11\n",
    "                         out_channels=2,\n",
    "                         channel_base=channel_base,\n",
    "                         activation=activation,\n",
    "                        )\n",
    "\n",
    "    def forward(self, x, scope):\n",
    "        inp = torch.cat((x, scope), 1)\n",
    "        logits = self.unet(inp)\n",
    "        alpha = torch.softmax(logits, 1)\n",
    "        # output channel 0 represents alpha_k,\n",
    "        # channel 1 represents (1 - alpha_k).\n",
    "        mask = scope * alpha[:, 0:1]\n",
    "        new_scope = scope * alpha[:, 1:2]\n",
    "        return mask, new_scope"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.2 Monet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderNet(nn.Module):\n",
    "    def __init__(self, in_channels, width, height, latent_size=16, mask_mode=\"concat\", activation=\"relu\"):\n",
    "        super().__init__()\n",
    "        self.mask_mode = mask_mode\n",
    "        self.convs = nn.Sequential(\n",
    "            nn.Conv2d(in_channels + 1 if mask_mode in [\"concat\", \"mulcat\"] else in_channels, 32, 3, stride=2),\n",
    "            get_activation(activation, inplace=True),\n",
    "            nn.Conv2d(32, 32, 3, stride=2),\n",
    "            get_activation(activation, inplace=True),\n",
    "            nn.Conv2d(32, 64, 3, stride=2),\n",
    "            get_activation(activation, inplace=True),\n",
    "            nn.Conv2d(64, 64, 3, stride=2),\n",
    "            get_activation(activation, inplace=True),\n",
    "        )\n",
    "#64x2x2, maybe small, maybe delete one layer.\n",
    "        for i in range(4):\n",
    "            width = (width - 1) // 2\n",
    "            height = (height - 1) // 2\n",
    "\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(64 * width * height, 256),\n",
    "            get_activation(activation, inplace=True),\n",
    "            nn.Linear(256, latent_size * 2)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.convs(x)\n",
    "        x = x.view(x.shape[0], -1)\n",
    "        x = self.mlp(x)\n",
    "        return x\n",
    "\n",
    "class DecoderNet(nn.Module):\n",
    "    def __init__(self, in_channels, height, width, latent_size, activation=\"relu\"):\n",
    "        super().__init__()\n",
    "        self.height = height\n",
    "        self.width = width\n",
    "        self.convs = nn.Sequential(\n",
    "            nn.Conv2d(latent_size + 2, 32, 3), # 18 = latent size() + 2\n",
    "            get_activation(activation, inplace=True),\n",
    "            nn.Conv2d(32, 32, 3),\n",
    "            get_activation(activation, inplace=True),\n",
    "            nn.Conv2d(32, 32, 3),\n",
    "            get_activation(activation, inplace=True),\n",
    "            nn.Conv2d(32, 32, 3),\n",
    "            get_activation(activation, inplace=True),\n",
    "            nn.Conv2d(32, in_channels + 1, 1),\n",
    "        )\n",
    "        ys = torch.linspace(-1, 1, self.height + 8) # 8 is for padding\n",
    "        xs = torch.linspace(-1, 1, self.width + 8)\n",
    "        ys, xs = torch.meshgrid(ys, xs)\n",
    "        coord_map = torch.stack((ys, xs)).unsqueeze(0)\n",
    "        self.register_buffer('coord_map_const', coord_map)\n",
    "\n",
    "    def forward(self, z):\n",
    "        z_tiled = z.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, self.height + 8, self.width + 8)\n",
    "        coord_map = self.coord_map_const.repeat(z.shape[0], 1, 1, 1)\n",
    "        inp = torch.cat((z_tiled, coord_map), 1)\n",
    "        result = self.convs(inp)\n",
    "        return result\n",
    "\n",
    "\n",
    "class Monet(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels,\n",
    "        height,\n",
    "        width,\n",
    "        n_blocks,\n",
    "        channel_base,\n",
    "        latent_size,\n",
    "        n_slots,\n",
    "        bg_sigma,\n",
    "        fg_sigma,\n",
    "        zero_color_weight=-1,\n",
    "        loss_type=\"gaussian\",\n",
    "        mask_mode=\"concat\",\n",
    "        activation=\"relu\",\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.in_channels = in_channels\n",
    "        self.height = height\n",
    "        self.width = width\n",
    "        self.n_blocks = n_blocks\n",
    "        self.channel_base = channel_base\n",
    "        self.latent_size = latent_size\n",
    "        self.n_slots = n_slots\n",
    "        self.bg_sigma = bg_sigma\n",
    "        self.fg_sigma = fg_sigma\n",
    "        self.loss_type = loss_type\n",
    "        self.mask_mode = mask_mode\n",
    "        self.activation = activation\n",
    "        self.zero_color_weight = zero_color_weight  # weight for the first channel. If -1, will use uniform\n",
    "        if self.zero_color_weight != -1:\n",
    "            self.weight = torch.cat([torch.tensor(zero_color_weight)[None], torch.ones(in_channels-1) * ((1-zero_color_weight)/(in_channels-1))])\n",
    "        self.attention = AttentionNet(\n",
    "            in_channels = in_channels,\n",
    "            n_blocks=n_blocks,\n",
    "            channel_base=channel_base,\n",
    "            activation=activation,\n",
    "        )\n",
    "        self.encoder = EncoderNet(\n",
    "            in_channels=in_channels,\n",
    "            width=width,\n",
    "            height=height,\n",
    "            latent_size=latent_size,\n",
    "            mask_mode=mask_mode,\n",
    "            activation=activation,\n",
    "        )\n",
    "        self.decoder = DecoderNet(\n",
    "            in_channels = in_channels,\n",
    "            width=width,\n",
    "            height=height,\n",
    "            latent_size=latent_size,\n",
    "            activation=activation,\n",
    "        )\n",
    "        self.beta = 0.5\n",
    "        self.gamma = 0.25\n",
    "\n",
    "    def forward(self, x):\n",
    "        scope = torch.ones_like(x[:, 0:1])\n",
    "        masks = []\n",
    "        for i in range(self.n_slots-1):\n",
    "            mask, scope = self.attention(x, scope)\n",
    "            masks.append(mask)\n",
    "        masks.append(scope)\n",
    "        loss = torch.zeros_like(x[:, 0, 0, 0])\n",
    "        mask_preds = []\n",
    "        full_reconstruction = torch.zeros_like(x)\n",
    "        if self.loss_type == \"gaussian\":\n",
    "            p_xs = torch.zeros_like(loss)\n",
    "        kl_zs = torch.zeros_like(loss)\n",
    "        for i, mask in enumerate(masks):\n",
    "            z, kl_z = self.__encoder_step(x, mask)\n",
    "            sigma = self.bg_sigma if i == 0 else self.fg_sigma\n",
    "            loss_ele, x_recon, mask_pred = self.__decoder_step(x, z, mask, sigma)\n",
    "            mask_preds.append(mask_pred)\n",
    "            loss += loss_ele + self.beta * kl_z\n",
    "            if self.loss_type == \"gaussian\":\n",
    "                p_xs += loss_ele\n",
    "            kl_zs += kl_z\n",
    "            full_reconstruction += mask * x_recon\n",
    "\n",
    "        masks = torch.cat(masks, 1)\n",
    "        tr_masks = torch.transpose(masks, 1, 3)\n",
    "        q_masks = dists.Categorical(probs=tr_masks)\n",
    "        q_masks_recon = dists.Categorical(logits=torch.stack(mask_preds, 3))\n",
    "        kl_masks = dists.kl_divergence(q_masks, q_masks_recon)\n",
    "        kl_masks = torch.sum(kl_masks, [1, 2])\n",
    "        # print('px', p_xs.mean().item(),\n",
    "        #       'kl_z', kl_zs.mean().item(),\n",
    "        #       'kl masks', kl_masks.mean().item())\n",
    "        loss += self.gamma * kl_masks\n",
    "\n",
    "        return {'loss': loss,\n",
    "                'masks': masks,\n",
    "                'reconstructions': full_reconstruction}\n",
    "\n",
    "\n",
    "    def __encoder_step(self, x, mask):\n",
    "        if self.mask_mode == \"concat\":\n",
    "            encoder_input = torch.cat((x, mask), 1)\n",
    "        elif self.mask_mode == \"mul\":\n",
    "            encoder_input = x * mask\n",
    "        elif self.mask_mode == \"mulcat\":\n",
    "            encoder_input = torch.cat((x*mask, mask), 1)\n",
    "        else:\n",
    "            raise\n",
    "        q_params = self.encoder(encoder_input)\n",
    "        # mlp output is latent size * 2\n",
    "\n",
    "        means = torch.sigmoid(q_params[:, :self.latent_size]) * 6 - 3\n",
    "        sigmas = torch.sigmoid(q_params[:, self.latent_size:]) * 3\n",
    "        dist = dists.Normal(means, sigmas)\n",
    "        dist_0 = dists.Normal(0., sigmas)\n",
    "        z = means + dist_0.sample()\n",
    "        q_z = dist.log_prob(z)\n",
    "        kl_z = dists.kl_divergence(dist, dists.Normal(0., 1.))\n",
    "        kl_z = torch.sum(kl_z, 1)\n",
    "        return z, kl_z\n",
    "\n",
    "    def __decoder_step(self, x, z, mask, sigma):\n",
    "        decoder_output = self.decoder(z)\n",
    "        mask_pred = decoder_output[:, self.in_channels]\n",
    "        if self.loss_type == \"gaussian\":\n",
    "            x_recon = torch.sigmoid(decoder_output[:, :self.in_channels]) # input channel (10)\n",
    "            dist = dists.Normal(x_recon, sigma)\n",
    "            p_x = dist.log_prob(x)\n",
    "            p_x *= mask\n",
    "            if self.zero_color_weight == -1:\n",
    "                p_x = torch.sum(p_x, [1, 2, 3])  #p_x: [B, C, H, W]\n",
    "            else:\n",
    "                p_x = torch.sum(p_x, [2, 3])\n",
    "                p_x = torch.matmul(p_x, self.weight.to(x.device))\n",
    "            loss_ele = -p_x\n",
    "        elif self.loss_type == \"ce\":\n",
    "            x_recon_logit = decoder_output[:, :self.in_channels]\n",
    "            if self.zero_color_weight == -1:\n",
    "                loss_ele = nn.CrossEntropyLoss(reduction='none')(x_recon_logit, x.argmax(1))\n",
    "            else:\n",
    "                loss_ele = nn.CrossEntropyLoss(reduction='none', weight=self.weight.to(x.device))(x_recon_logit, x.argmax(1))\n",
    "            loss_ele *= mask.squeeze(1)  # mask: [B, 1, H, W], loss_ele: [B, H, W]\n",
    "            loss_ele = torch.sum(loss_ele, [-2,-1])\n",
    "            x_recon = nn.Softmax(dim=1)(x_recon_logit)\n",
    "        elif self.loss_type == \"gaussian+ce\":\n",
    "            x_recon_logit = decoder_output[:, :self.in_channels]\n",
    "            x_recon = nn.Softmax(dim=1)(x_recon_logit)\n",
    "\n",
    "            # gaussian:\n",
    "            dist = dists.Normal(x_recon, sigma)\n",
    "            p_x = dist.log_prob(x)\n",
    "            p_x *= mask\n",
    "            if self.zero_color_weight == -1:\n",
    "                p_x = torch.sum(p_x, [1, 2, 3])  #p_x: [B, C, H, W]\n",
    "            else:\n",
    "                p_x = torch.sum(p_x, [2, 3])\n",
    "                p_x = torch.matmul(p_x, self.weight.to(x.device))\n",
    "            loss_ele_gaussian = -p_x\n",
    "\n",
    "            # ce:\n",
    "            if self.zero_color_weight == -1:\n",
    "                loss_ele_ce = nn.CrossEntropyLoss(reduction='none')(x_recon_logit, x.argmax(1))\n",
    "            else:\n",
    "                loss_ele_ce = nn.CrossEntropyLoss(reduction='none', weight=self.weight.to(x.device))(x_recon_logit, x.argmax(1))\n",
    "            loss_ele_ce *= mask.squeeze(1)\n",
    "            loss_ele_ce = torch.sum(loss_ele_ce, [-2,-1])\n",
    "\n",
    "            loss_ele = loss_ele_gaussian + loss_ele_ce\n",
    "        else:\n",
    "            raise Exception(\"loss_type '{}' is not valid!\".format(self.loss_type))\n",
    "        return loss_ele, x_recon, mask_pred\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Monet\"}\n",
    "        model_dict[\"in_channels\"] = self.in_channels\n",
    "        model_dict[\"height\"] = self.height\n",
    "        model_dict[\"width\"] = self.width\n",
    "        model_dict[\"n_blocks\"] = self.n_blocks\n",
    "        model_dict[\"channel_base\"] = self.channel_base\n",
    "        model_dict[\"latent_size\"] = self.latent_size\n",
    "        model_dict[\"n_slots\"] = self.n_slots\n",
    "        model_dict[\"bg_sigma\"] = self.bg_sigma\n",
    "        model_dict[\"fg_sigma\"] = self.fg_sigma\n",
    "        model_dict[\"mask_mode\"] = self.mask_mode\n",
    "        model_dict[\"loss_type\"] = self.loss_type\n",
    "        model_dict[\"activation\"] = self.activation\n",
    "        model_dict[\"zero_color_weight\"] = self.zero_color_weight\n",
    "        model_dict[\"state_dict\"] = to_cpu(self.state_dict())\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "def load_model(model_dict, device=\"cpu\"):\n",
    "    model_type = model_dict[\"type\"]\n",
    "    if model_type == \"Monet\":\n",
    "        model = Monet(\n",
    "            in_channels=model_dict[\"in_channels\"],\n",
    "            height=model_dict[\"height\"],\n",
    "            width=model_dict[\"width\"],\n",
    "            n_blocks=model_dict[\"n_blocks\"],\n",
    "            channel_base=model_dict[\"channel_base\"],\n",
    "            latent_size=model_dict[\"latent_size\"],\n",
    "            n_slots=model_dict[\"n_slots\"],\n",
    "            bg_sigma=model_dict[\"bg_sigma\"],\n",
    "            fg_sigma=model_dict[\"fg_sigma\"],\n",
    "            loss_type=model_dict[\"loss_type\"] if \"loss_type\" in model_dict else \"gaussian\",\n",
    "            mask_mode=model_dict[\"mask_mode\"] if \"mask_mode\" in model_dict else \"concat\",\n",
    "            activation=model_dict[\"activation\"] if \"activation\" in model_dict else \"relu\",\n",
    "            zero_color_weight=model_dict[\"zero_color_weight\"] if \"zero_color_weight\" in model_dict else -1,\n",
    "        )\n",
    "    else:\n",
    "        raise\n",
    "    model.load_state_dict(model_dict[\"state_dict\"])\n",
    "    model.to(device)\n",
    "    return model"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}