3.TrainStudentHint part 1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1a865bf-d19f-4803-b256-4a3247bd4667",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from time import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "\n",
    "\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from scipy.sparse.linalg import eigs\n",
    "\n",
    "\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "DEVICE = torch.device('cuda:0')\n",
    "print(\"CUDA:\", USE_CUDA, DEVICE)\n",
    "\n",
    "from tensorboardX import SummaryWriter\n",
    "# sw = SummaryWriter(logdir='./Student_HINT', flush_secs=5)\n",
    "\n",
    "import math\n",
    "from typing import Optional, List, Union\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import Parameter\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.typing import OptTensor\n",
    "from torch_geometric.nn.conv import MessagePassing\n",
    "from torch_geometric.transforms import LaplacianLambdaMax\n",
    "from torch_geometric.utils import remove_self_loops, add_self_loops, get_laplacian\n",
    "from torch_geometric.utils import to_dense_adj\n",
    "from torch_scatter import scatter_add\n",
    "# from .autonotebook import tqdm as notebook_tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53d26d62-aed2-4d96-bc93-9efda7479480",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_graphdata_channel1(graph_signal_matrix_filename, num_of_hours, num_of_days, num_of_weeks, batch_size,\n",
    "                            shuffle=True, DEVICE = torch.device('cuda:0')):\n",
    "    '''\n",
    "    :param graph_signal_matrix_filename: str\n",
    "    :param num_of_hours: int\n",
    "    :param num_of_days: int\n",
    "    :param num_of_weeks: int\n",
    "    :param DEVICE:\n",
    "    :param batch_size: int\n",
    "    :return:\n",
    "    three DataLoaders, each dataloader contains:\n",
    "    test_x_tensor: (B, N_nodes, in_feature, T_input)\n",
    "    test_decoder_input_tensor: (B, N_nodes, T_output)\n",
    "    test_target_tensor: (B, N_nodes, T_output)\n",
    "    '''\n",
    "\n",
    "    file = os.path.basename(graph_signal_matrix_filename).split('.')[0]\n",
    "    filename = os.path.join('./data/PEMS04/', file + '_r' + str(num_of_hours) + '_d' + str(num_of_days) + '_w' + str(num_of_weeks)) +'_astcgn'\n",
    "    print('load file:', filename)\n",
    "\n",
    "    file_data = np.load(filename + '.npz')\n",
    "    train_x = file_data['train_x']  # (10181, 307, 3, 12)\n",
    "    train_x = train_x[:, :, 0:1, :]\n",
    "    train_target = file_data['train_target']  # (10181, 307, 12)\n",
    "\n",
    "    val_x = file_data['val_x']\n",
    "    val_x = val_x[:, :, 0:1, :]\n",
    "    val_target = file_data['val_target']\n",
    "\n",
    "    test_x = file_data['test_x']\n",
    "    test_x = test_x[:, :, 0:1, :]\n",
    "    test_target = file_data['test_target']\n",
    "\n",
    "    mean = file_data['mean'][:, :, 0:1, :]  # (1, 1, 3, 1)\n",
    "    std = file_data['std'][:, :, 0:1, :]  # (1, 1, 3, 1)\n",
    "\n",
    "    # ------- train_loader -------\n",
    "    train_x_tensor = torch.from_numpy(train_x).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)\n",
    "    train_target_tensor = torch.from_numpy(train_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)\n",
    "    train_dataset = torch.utils.data.TensorDataset(train_x_tensor, train_target_tensor)\n",
    "    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle)\n",
    "\n",
    "    # ------- val_loader -------\n",
    "    val_x_tensor = torch.from_numpy(val_x).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)\n",
    "    val_target_tensor = torch.from_numpy(val_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)\n",
    "    val_dataset = torch.utils.data.TensorDataset(val_x_tensor, val_target_tensor)\n",
    "    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    # ------- test_loader -------\n",
    "    test_x_tensor = torch.from_numpy(test_x).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)\n",
    "    test_target_tensor = torch.from_numpy(test_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)\n",
    "    test_dataset = torch.utils.data.TensorDataset(test_x_tensor, test_target_tensor)\n",
    "    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    # print\n",
    "    print('train:', train_x_tensor.size(), train_target_tensor.size())\n",
    "    print('val:', val_x_tensor.size(), val_target_tensor.size())\n",
    "    print('test:', test_x_tensor.size(), test_target_tensor.size())\n",
    "\n",
    "    return train_loader, train_target_tensor, val_loader, val_target_tensor, test_loader, test_target_tensor, mean, std\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3913859a-5cc7-4103-803c-c01c07cc624b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "graph_signal_matrix_filename = './data/PEMS04/pems04.npz'\n",
    "batch_size = 16\n",
    "num_of_weeks = 1\n",
    "num_of_days = 1\n",
    "num_of_hours = 1\n",
    "\n",
    "train_loader, train_target_tensor, val_loader, val_target_tensor, test_loader, test_target_tensor, _mean, _std = load_graphdata_channel1(\n",
    "    graph_signal_matrix_filename, num_of_hours, num_of_days, num_of_weeks, batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eba75207-cb40-4e13-b240-b5db01af551c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_adjacency_matrix(distance_df_filename, num_of_vertices, id_filename=None):\n",
    "    '''\n",
    "    Parameters\n",
    "    ----------\n",
    "    distance_df_filename: str, path of the csv file contains edges information\n",
    "    num_of_vertices: int, the number of vertices\n",
    "    Returns\n",
    "    ----------\n",
    "    A: np.ndarray, adjacency matrix\n",
    "    '''\n",
    "    if 'npy' in distance_df_filename:  # false\n",
    "        adj_mx = np.load(distance_df_filename)\n",
    "        return adj_mx, None\n",
    "    else:\n",
    "        \n",
    "        #--------------------------------------------- read from here\n",
    "        import csv\n",
    "        A = np.zeros((int(num_of_vertices), int(num_of_vertices)),dtype=np.float32)\n",
    "        distaneA = np.zeros((int(num_of_vertices), int(num_of_vertices)), dtype=np.float32)\n",
    "\n",
    "        #------------ Ignore\n",
    "        if id_filename: # false\n",
    "            with open(id_filename, 'r') as f:\n",
    "                id_dict = {int(i): idx for idx, i in enumerate(f.read().strip().split('\\n'))}  # 把节点id（idx）映射成从0开始的索引\n",
    "\n",
    "            with open(distance_df_filename, 'r') as f:\n",
    "                f.readline()\n",
    "                reader = csv.reader(f)\n",
    "                for row in reader:\n",
    "                    if len(row) != 3:\n",
    "                        continue\n",
    "                    i, j, distance = int(row[0]), int(row[1]), float(row[2])\n",
    "                    A[id_dict[i], id_dict[j]] = 1\n",
    "                    distaneA[id_dict[i], id_dict[j]] = distance\n",
    "            return A, distaneA\n",
    "\n",
    "        else:\n",
    "         #-------------Continue reading\n",
    "            with open(distance_df_filename, 'r') as f:\n",
    "                f.readline()\n",
    "                reader = csv.reader(f)\n",
    "                for row in reader:\n",
    "                    if len(row) != 3:\n",
    "                        continue\n",
    "                    i, j, distance = int(row[0]), int(row[1]), float(row[2])\n",
    "                    A[i, j] = 1\n",
    "                    distaneA[i, j] = distance\n",
    "            return A, distaneA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "772db769-0b01-4e86-980d-249a12cd6b72",
   "metadata": {},
   "outputs": [],
   "source": [
    "id_filename = None\n",
    "adj_filename = './data/PEMS04/PEMS04.csv'\n",
    "num_of_vertices = 307\n",
    "adj_mx, distance_mx = get_adjacency_matrix(adj_filename, num_of_vertices, id_filename) #  adj_mx and distance_mx (307, 307)\n",
    "\n",
    "rows, cols = np.where(adj_mx == 1)\n",
    "edges = zip(rows.tolist(), cols.tolist())\n",
    "gr = nx.Graph()\n",
    "gr.add_edges_from(edges)\n",
    "# nx.draw(gr, node_size=3)\n",
    "plt.show()\n",
    "rows, cols = np.where(adj_mx == 1)\n",
    "edges = zip(rows.tolist(), cols.tolist())\n",
    "edge_index_data = torch.LongTensor(np.array([rows, cols])).to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9d96052-6506-427c-85fe-a62e35748109",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.ASTGCN import ASTGCN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebcc7a3f-54c0-4732-9320-a48a0a960f13",
   "metadata": {},
   "outputs": [],
   "source": [
    "nb_block = 2\n",
    "in_channels = 1\n",
    "K = 3\n",
    "nb_chev_filter =64\n",
    "nb_time_filter = 64\n",
    "time_strides = num_of_hours\n",
    "num_for_predict = 12\n",
    "len_input = 12\n",
    "dropout=0.01\n",
    "\n",
    "#L_tilde = scaled_Laplacian(adj_mx)\n",
    "#cheb_polynomials = [torch.from_numpy(i).type(torch.FloatTensor).to(DEVICE) for i in cheb_polynomial(L_tilde, K)]\n",
    "# net = ASTGCN( nb_block, in_channels, K, nb_chev_filter, nb_time_filter, time_strides, num_for_predict, len_input, num_of_vertices).to(DEVICE)\n",
    "Teacher=ASTGCN( nb_block, in_channels, K, nb_chev_filter, nb_time_filter, time_strides, num_for_predict, len_input, num_of_vertices,dropout).to(DEVICE)\n",
    "print(Teacher)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50dee06c-5f15-457b-80bf-e9538b7792a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "nb_chev_filterstdn = 64\n",
    "nb_time_filterstdn = 64\n",
    "nb_blockstdn = 1\n",
    "dropouts=0.01\n",
    "Student=ASTGCN( nb_blockstdn, in_channels, K, nb_chev_filterstdn, nb_time_filterstdn, time_strides, num_for_predict, len_input, num_of_vertices,dropouts).to(DEVICE)\n",
    "print(Student)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0da0e88-db3d-4690-baaa-2d5f1939e756",
   "metadata": {},
   "source": [
    "# intialized teacher model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d119e88c-aab0-4bbf-84bd-f07093dd3974",
   "metadata": {},
   "outputs": [],
   "source": [
    "#------------------------------------------------------\n",
    "learning_rate = 0.001\n",
    "# \n",
    "optimizerTeacher = optim.Adam(Teacher.parameters(), lr=1e-3,weight_decay=1e-6)#l2 regularization applied in weightdecay<0\n",
    "\n",
    "print('Teacher\\'s state_dict:')\n",
    "total_param = 0\n",
    "for param_tensor in Teacher.state_dict():\n",
    "    print(param_tensor, '\\t', Teacher.state_dict()[param_tensor].size(), '\\t', Teacher.state_dict()[param_tensor].device)\n",
    "    total_param += np.prod(Teacher.state_dict()[param_tensor].size())\n",
    "print('Teacher\\'s total params:', total_param)\n",
    "#--------------------------------------------------\n",
    "print('Optimizer\\'s state_dict:')\n",
    "for var_name in optimizerTeacher.state_dict():\n",
    "    print(var_name, '\\t', optimizerTeacher.state_dict()[var_name])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fb1180f-7517-4d87-9982-dbd7c64860a9",
   "metadata": {},
   "source": [
    "## Intialized student model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2ae6e1d-3714-4695-bbae-16112488590e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#------------------------------------------------------\n",
    "# weight_decay=1e-3\n",
    "optimizerStudent = optim.Adam(Student.parameters(), lr=1e-3,weight_decay=1e-6)\n",
    "\n",
    "print('Teacher\\'s state_dict:')\n",
    "total_param = 0\n",
    "for param_tensor in Student.state_dict():\n",
    "    print(param_tensor, '\\t', Student.state_dict()[param_tensor].size(), '\\t', Student.state_dict()[param_tensor].device)\n",
    "    total_param += np.prod(Student.state_dict()[param_tensor].size())\n",
    "print('Student\\'s total params:', total_param)\n",
    "#--------------------------------------------------\n",
    "print('Optimizer\\'s state_dict:')\n",
    "for var_name in optimizerStudent.state_dict():\n",
    "    print(var_name, '\\t', optimizerStudent.state_dict()[var_name])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99cdfc51-7445-4886-9dd8-78649487c0d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def masked_mae(preds, labels, null_val=np.nan):\n",
    "    if np.isnan(null_val):\n",
    "        mask = ~torch.isnan(labels)\n",
    "    else:\n",
    "        mask = (labels != null_val)\n",
    "    mask = mask.float()\n",
    "    mask /= torch.mean((mask))\n",
    "    mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask)\n",
    "    loss = torch.abs(preds - labels)\n",
    "    loss = loss * mask\n",
    "    loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss)\n",
    "    return torch.mean(loss)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "778ff538-0c64-4191-9752-f5fe61f5d047",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RMSELoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.mse = nn.MSELoss()\n",
    "        \n",
    "    def forward(self,ytil,y):\n",
    "        return torch.sqrt(self.mse(ytil,y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "142d4546-b1a4-41fa-9e86-7056d8016980",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha=0.4# trade off teacher weights to students, 70% teach 30% stud\n",
    "masked_flag=0\n",
    "criterionStudent=RMSELoss().to(DEVICE)\n",
    "\n",
    "criterion_masked = masked_mae\n",
    "loss_function = 'mae'\n",
    "\n",
    "metric_method = 'unmask'\n",
    "missing_value=0.0\n",
    "\n",
    "if loss_function=='masked_mse':\n",
    "    criterion_masked = masked_mse         #nn.MSELoss().to(DEVICE)\n",
    "    masked_flag=1\n",
    "elif loss_function=='masked_mae':\n",
    "    criterion_masked = masked_mae\n",
    "    masked_flag = 1\n",
    "elif loss_function == 'mae':\n",
    "    criterion = nn.L1Loss().to(DEVICE)\n",
    "    masked_flag = 0\n",
    "elif loss_function == 'rmse':\n",
    "    criterion = nn.MSELoss().to(DEVICE)\n",
    "    masked_flag= 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0fea845-a7da-4200-a5ad-f4e409f4fe0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ValLoss=[]\n",
    "TrainLoss=[]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06cc7bee-ed67-447c-b642-27303b61a7fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_val_loss_mstgcn(net, val_loader, criterion,  masked_flag,missing_value, epoch, edge_index_data, limit=None):\n",
    "    '''\n",
    "    for rnn, compute mean loss on validation set\n",
    "    :param net: model\n",
    "    :param val_loader: torch.utils.data.utils.DataLoader\n",
    "    :param criterion: torch.nn.MSELoss\n",
    "    :param sw: tensorboardX.SummaryWriter\n",
    "    :param global_step: int, current global_step\n",
    "    :param limit: int,\n",
    "    :return: val_loss\n",
    "    '''\n",
    "    net.train(False)  # ensure dropout layers are in evaluation mode\n",
    "    with torch.no_grad():\n",
    "        val_loader_length = len(val_loader)  # nb of batch\n",
    "        tmp = []  # batch loss\n",
    "        for batch_index, batch_data in enumerate(val_loader):\n",
    "            encoder_inputs, labels = batch_data\n",
    "            outputs = net(encoder_inputs, edge_index_data)\n",
    "            if masked_flag:\n",
    "                loss = criterion(outputs, labels)\n",
    "            else:\n",
    "                loss = criterion(outputs, labels)\n",
    "            tmp.append(loss.item())\n",
    "            if batch_index % 20 == 0:\n",
    "                print('validation batch %s / %s, loss: %.2f' % (batch_index + 1, val_loader_length, loss.item()))\n",
    "            if (limit is not None) and batch_index >= limit:\n",
    "                break\n",
    "\n",
    "        validation_loss = sum(tmp) / len(tmp)\n",
    "        ValLoss.append(validation_loss)\n",
    "        # sw.add_scalar('validation_loss', validation_loss, epoch)\n",
    "        print(\"rata rata val_loss\",validation_loss,\"pada epoch:\",epoch)\n",
    "    return validation_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24a1642d-f5b7-4524-9ea5-a3856b045e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "global_step = 0\n",
    "best_epoch = 0\n",
    "best_val_loss = np.inf\n",
    "start_time= time()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f18b437e-8379-420d-a6a8-dea3f043807b",
   "metadata": {},
   "source": [
    "# load weights from teacher"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deb20d13-5500-4c21-8a83-371607e2b86e",
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch =  22\n",
    "params_filename = os.path.join('./Teacher/', 'checkpoint_%s.pth' % epoch)\n",
    "loaded_checkpoint=torch.load(params_filename)\n",
    "epoch=loaded_checkpoint['epoch']\n",
    "Teacher.load_state_dict(loaded_checkpoint['model_state'])\n",
    "optimizerTeacher.load_state_dict(loaded_checkpoint['optimizer_state'])\n",
    "Teacher.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9efb359b-392a-48fd-8262-0cba113e44dc",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(Teacher.state_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cf41d60-b8b2-4b5a-a4e2-11b433d6a8e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "a=os.path.join('./Teacher/','N.npy')\n",
    "N=np.load(a)\n",
    "N=torch.tensor(N)\n",
    "N.to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "711a8e5d-779f-4b13-88f8-9a28dade0a3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(Teacher)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fd454de-5805-4085-8813-5f3e728b4e1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "a=torch.nn.functional.mse_loss\n",
    "a(Student._blocklist[0]._residual_convolution.weight,Teacher._blocklist[0]._residual_convolution.weight)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5de9be17-2295-41c2-8d3c-ddafe675f17c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChebConvAttention(MessagePassing):\n",
    "\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels: int,\n",
    "        out_channels: int,\n",
    "        K: int,\n",
    "        normalization: Optional[str] = None,\n",
    "        bias: bool = True,\n",
    "        **kwargs\n",
    "    ):\n",
    "        kwargs.setdefault(\"aggr\", \"add\")\n",
    "        super(ChebConvAttention, self).__init__(**kwargs)\n",
    "\n",
    "        assert K > 0\n",
    "        assert normalization in [None, \"sym\", \"rw\"], \"Invalid normalization\"\n",
    "\n",
    "        self._in_channels = in_channels\n",
    "        self._out_channels = out_channels\n",
    "        self._normalization = normalization\n",
    "        self._weight = Parameter(torch.Tensor(K, in_channels, out_channels)) \n",
    "\n",
    "        if bias:\n",
    "            self._bias = Parameter(torch.Tensor(out_channels))\n",
    "        else:\n",
    "            self.register_parameter(\"_bias\", None)\n",
    "\n",
    "        self._reset_parameters()\n",
    "\n",
    "    def _reset_parameters(self):\n",
    "        nn.init.xavier_uniform_(self._weight)\n",
    "        if self._bias is not None:\n",
    "            nn.init.uniform_(self._bias)\n",
    "\n",
    "    #--forward pass-----\n",
    "    def __norm__(\n",
    "        self,\n",
    "        edge_index,\n",
    "        num_nodes: Optional[int],\n",
    "        edge_weight: OptTensor,\n",
    "        normalization: Optional[str],\n",
    "        lambda_max,\n",
    "        dtype: Optional[int] = None,\n",
    "        batch: OptTensor = None,\n",
    "    ):\n",
    "\n",
    "        edge_index, edge_weight = remove_self_loops(edge_index, edge_weight)\n",
    "\n",
    "        edge_index, edge_weight = get_laplacian(\n",
    "            edge_index, edge_weight, normalization, dtype, num_nodes\n",
    "        )\n",
    "\n",
    "        if batch is not None and lambda_max.numel() > 1:\n",
    "            lambda_max = lambda_max[batch[edge_index[0]]]\n",
    "\n",
    "        edge_weight = (2.0 * edge_weight) / lambda_max\n",
    "        edge_weight.masked_fill_(edge_weight == float(\"inf\"), 0)\n",
    "\n",
    "        edge_index, edge_weight = add_self_loops(\n",
    "            edge_index, edge_weight, fill_value=-1.0, num_nodes=num_nodes\n",
    "        )\n",
    "        assert edge_weight is not None\n",
    "\n",
    "        return edge_index, edge_weight #for example 307 nodes as deg, 340 edges , 307 nodes as self connections\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        x: torch.FloatTensor,\n",
    "        edge_index: torch.LongTensor,\n",
    "        spatial_attention: torch.FloatTensor,\n",
    "        edge_weight: OptTensor = None,\n",
    "        batch: OptTensor = None,\n",
    "        lambda_max: OptTensor = None,\n",
    "    ) -> torch.FloatTensor:\n",
    "\n",
    "        if self._normalization != \"sym\" and lambda_max is None:\n",
    "            raise ValueError(\n",
    "                \"You need to pass `lambda_max` to `forward() in`\"\n",
    "                \"case the normalization is non-symmetric.\"\n",
    "            )\n",
    "\n",
    "        if lambda_max is None:\n",
    "            lambda_max = torch.tensor(2.0, dtype=x.dtype, device=x.device)\n",
    "        if not isinstance(lambda_max, torch.Tensor):\n",
    "            lambda_max = torch.tensor(lambda_max, dtype=x.dtype, device=x.device)\n",
    "        assert lambda_max is not None\n",
    "\n",
    "        edge_index, norm = self.__norm__(\n",
    "            edge_index,\n",
    "            x.size(self.node_dim),\n",
    "            edge_weight,\n",
    "            self._normalization,\n",
    "            lambda_max,\n",
    "            dtype=x.dtype,\n",
    "            batch=batch,\n",
    "        )\n",
    "        row, col = edge_index # refer to the index of each note each is a list of nodes not a number # (954, 954)\n",
    "        Att_norm = norm * spatial_attention[:, row, col] # spatial_attention for example (32, 307, 307), -> (954) * (32, 954) -> (32, 954)\n",
    "        num_nodes = x.size(self.node_dim) #for example 307\n",
    "        # (307, 307) * (32, 307, 307) -> (32, 307, 307) -permute-> (32, 307,307) * (32, 307, 1) -> (32, 307, 1)\n",
    "        TAx_0 = torch.matmul(\n",
    "            (torch.eye(num_nodes).to(edge_index.device) * spatial_attention).permute(\n",
    "                0, 2, 1\n",
    "            ),\n",
    "            x,\n",
    "        ) #for example (32, 307, 1)\n",
    "        out = torch.matmul(TAx_0, self._weight[0]) #for example (32, 307, 1) * [1, 64] -> (32, 307, 64)\n",
    "        edge_index_transpose = edge_index[[1, 0]]\n",
    "        if self._weight.size(0) > 1:\n",
    "            TAx_1 = self.propagate(\n",
    "                edge_index_transpose, x=TAx_0, norm=Att_norm, size=None\n",
    "            )\n",
    "            out = out + torch.matmul(TAx_1, self._weight[1])\n",
    "\n",
    "        for k in range(2, self._weight.size(0)):\n",
    "            TAx_2 = self.propagate(edge_index_transpose, x=TAx_1, norm=norm, size=None)\n",
    "            TAx_2 = 2.0 * TAx_2 - TAx_0\n",
    "            out = out + torch.matmul(TAx_2, self._weight[k])\n",
    "            TAx_0, TAx_1 = TAx_1, TAx_2\n",
    "\n",
    "        if self._bias is not None:\n",
    "            out += self._bias\n",
    "\n",
    "        return out #? (b, N, F_out) (32, 307, 64)\n",
    "\n",
    "\n",
    "    def message(self, x_j, norm):\n",
    "        if norm.dim() == 1:  # true\n",
    "            return norm.view(-1, 1) * x_j  # (954, 1) * (32, 954, 1) -> (32, 954, 1)\n",
    "        else:\n",
    "            d1, d2 = norm.shape\n",
    "            return norm.view(d1, d2, 1) * x_j\n",
    "\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"{}({}, {}, K={}, normalization={})\".format(\n",
    "            self.__class__.__name__,\n",
    "            self._in_channels,\n",
    "            self._out_channels,\n",
    "            self._weight.size(0),\n",
    "            self._normalization,\n",
    "        )\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de1f6a18-8eae-43c3-a730-8ede1d8e25a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "W_gcheb=Student._blocklist[0]._chebconv_attention._weight\n",
    "W_cheb2=Teacher._blocklist[1]._chebconv_attention._weight\n",
    "\n",
    "W_gtime=Student._blocklist[0]._time_convolution.weight\n",
    "W_time2=Teacher._blocklist[1]._time_convolution.weight\n",
    "\n",
    "W_cheb2.requires_grad = False\n",
    "W_time2.requires_grad = False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a58422da-c6fe-4dba-8c3e-713a83b1b6bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def W_regcheb\n",
    "regcheb=ChebConvAttention(63, 64, K=3, normalization=None)\n",
    "W_regcheb=regcheb._weight.to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0e7469e-e04c-400b-9f98-7fcd33f2383b",
   "metadata": {},
   "outputs": [],
   "source": [
    "W_regcheb.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c02acc7-fbf1-49fa-990e-4deaf2f25483",
   "metadata": {},
   "outputs": [],
   "source": [
    "# W_regtime.shape()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3b85c7b-a993-415e-a446-3077c2f61ebf",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print('Teacher\\'s state_dict:')\n",
    "epoch = best_epoch\n",
    "total_param \n",
    "for param_tensor in Teacher.state_dict():\n",
    "    print(param_tensor, '\\t', Teacher.state_dict()[param_tensor].size(), '\\t', Teacher.state_dict()[param_tensor].device)\n",
    "    total_param += np.prod(Teacher.state_dict()[param_tensor].size())\n",
    "print('Teacher\\'s total params:', total_param)\n",
    "#--------------------------------------------------\n",
    "print('Optimizer\\'s state_dict:')\n",
    "for var_name in optimizerTeacher.state_dict():\n",
    "    print(var_name, '\\t', optimizerTeacher.state_dict()[var_name])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "776bffa9-b938-40fd-8a98-69c8048d8078",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print('Teacher\\'s state_dict:')\n",
    "total_param = 0\n",
    "for param_tensor in Student.state_dict():\n",
    "    print(param_tensor, '\\t', Student.state_dict()[param_tensor].size(), '\\t', Student.state_dict()[param_tensor].device)\n",
    "    total_param += np.prod(Student.state_dict()[param_tensor].size())\n",
    "print('Student\\'s total params:', total_param)\n",
    "#--------------------------------------------------\n",
    "print('Optimizer\\'s state_dict:')\n",
    "for var_name in optimizerStudent.state_dict():\n",
    "    print(var_name, '\\t', optimizerStudent.state_dict()[var_name])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef55aa0c-f601-408f-a045-e157930b11d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_layer = np.inf\n",
    "best_layer2 = np.inf\n",
    "best_epoch=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55d3dbee-322d-49c5-91cf-9a3b131012bd",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "g=0\n",
    "layer2=[]\n",
    "layer3=[]\n",
    "for epoch in range(40):\n",
    "    generalization= os.path.join('./Student_HINT/', 'Generalized+tc+cc_%s.pth' % epoch)\n",
    "    Student.train()  # ensure dropout layers are in train mode\n",
    "    Teacher.eval()\n",
    "    t=[]\n",
    "    c=[]\n",
    "   \n",
    "    for batch_index, batch_data in enumerate(train_loader):\n",
    "        encoder_inputs, labels = batch_data   # encoder_inputs torch.Size([32, 307, 1, 12])  label torch.Size([32, 307, 12])\n",
    "        P_s = Student(encoder_inputs, edge_index_data) # torch.Size([32, 307, 12]) Ps\n",
    "        with torch.no_grad():\n",
    "            P_t=Teacher(encoder_inputs, edge_index_data)# Pt\n",
    "        P_tloss=criterionStudent(P_t,labels)\n",
    "        P_sloss=criterionStudent(P_s,labels)\n",
    "        phi=(1-(P_tloss/N))\n",
    "        #size (cheb2 3, 64, 64)\n",
    "        #size (gcheb1 3, 1, 64)\n",
    "        #size (regcheb 3, 63, 64)\n",
    "        optimizerStudent.zero_grad()\n",
    "        losst=phi*criterionStudent(W_gtime,W_time2)\n",
    "        losst.backward()\n",
    "        t.append(losst.item())\n",
    "        \n",
    "        lossc=phi*criterionStudent(W_gcheb[:3,:1,:64],W_cheb2[:3,:1,:64])\n",
    "        lossc1=phi*criterionStudent(W_regcheb[:3,0:63,:64],W_cheb2[:3,1:64,:64])\n",
    "        lossct=+lossc+lossc1\n",
    "        lossct.backward()\n",
    "        c.append(lossct.item())  \n",
    "        optimizerStudent.step()\n",
    "        g+=1 \n",
    "        if g%800 == 0:\n",
    "            print('step', g)\n",
    "    ts=sum(t)/len(t)\n",
    "    cs=sum(c)/len(c)\n",
    "    \n",
    "    layer2.append(ts)\n",
    "    layer3.append(cs)\n",
    "    print('loss per layer pada epoch', epoch)\n",
    "    if ts < best_layer:\n",
    "        best_layer = ts\n",
    "        epoch=best_epoch\n",
    "        print('#### sucessfuly saving model ! ####')\n",
    "        checkpoint = {\n",
    "            'model_state': Student.state_dict(),\n",
    "            'optimizer_state': optimizerStudent.state_dict()\n",
    "        }\n",
    "        torch.save(checkpoint, generalization)\n",
    "    print('timeconv:',ts)\n",
    "    print('chebyconv:',cs)\n",
    "   \n",
    "    \n",
    "    \n",
    "\n",
    "path1=os.path.join(\"./Student_HINT/\",'timeconv.npy')\n",
    "path2=os.path.join(\"./Student_HINT/\",'chebychev.npy')\n",
    "\n",
    "np.save(path1,layer2)\n",
    "np.save(path2,layer3)\n",
    "print(best_epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a19667d-5e14-4b95-9352-b85712442d75",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = {\n",
    "            'model_state': Student.state_dict(),\n",
    "            'optimizer_state': optimizerStudent.state_dict()\n",
    "        }\n",
    "torch.save(checkpoint, generalization)"
   ]
  }
 ],
 "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}