{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%matplotlib inline" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Model Fusion by Graph Matching\n\nThis example shows how to fuse different models into a single model by ``pygmtools``.\nModel fusion aims to fuse multiple models into one, such that the fused model could have higher performance.\nThe neural networks can be regarded as graphs (channels - nodes, update functions between channels - edges;\nnode feature - bias, edge feature - weights), and fusing the models is equivalent to solving a graph matching\nproblem. In this example, the given models are trained on MNIST data from different distributions, and the\nfused model could combine the knowledge from two input models and can reach higher accuracy when testing.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Author: Chang Liu \n# Runzhong Wang \n#\n# License: Mulan PSL v2 License" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "

Note

This is a simplified implementation of the ideas in [Liu et al. Deep Neural Network Fusion via Graph Matching with Applications to Model Ensemble and Federated Learning. ICML 2022.](https://proceedings.mlr.press/v162/liu22k/liu22k.pdf)\n For more details, please refer to the paper and the [official code repository](https://github.com/Thinklab-SJTU/GAMF).

\n\n

Note

The following solvers are included in this example:\n\n * :func:`~pygmtools.classic_solvers.sm` (classic solver)\n\n * :func:`~pygmtools.linear_solvers.hungarian` (linear solver)

\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nimport torchvision.transforms as transforms\nimport time\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nimport pygmtools as pygm\n\npygm.BACKEND = 'pytorch'\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Define a simple CNN classifier network\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "class SimpleNet(nn.Module):\n def __init__(self):\n super(SimpleNet, self).__init__()\n self.conv1 = nn.Conv2d(1, 32, 5, padding=1, padding_mode='replicate', bias=False)\n self.max_pool = nn.MaxPool2d(2, padding=1)\n self.conv2 = nn.Conv2d(32, 64, 5, padding=1, padding_mode='replicate', bias=False)\n self.fc1 = nn.Linear(3136, 32, bias=False)\n self.fc2 = nn.Linear(32, 10, bias=False)\n\n def forward(self, x):\n output = F.relu(self.conv1(x))\n output = self.max_pool(output)\n output = F.relu(self.conv2(output))\n output = self.max_pool(output)\n output = output.view(output.shape[0], -1)\n output = self.fc1(output)\n output = self.fc2(output)\n return output" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Load the trained models to be fused\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "model1 = SimpleNet()\nmodel2 = SimpleNet()\nmodel1.load_state_dict(torch.load('../data/example_model_fusion_1.dat', map_location=device))\nmodel2.load_state_dict(torch.load('../data/example_model_fusion_2.dat', map_location=device))\nmodel1.to(device)\nmodel2.to(device)\ntest_dataset = torchvision.datasets.MNIST(\n root='../data/mnist_data', # the directory to store the dataset\n train=False, # the dataset is used to test\n transform=transforms.ToTensor(), # the dataset is in the form of tensors\n download=True)\ntest_loader = torch.utils.data.DataLoader(\n dataset=test_dataset,\n batch_size=32,\n shuffle=False)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Print the layers of the simple CNN model:\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "print(model1)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Test the input models\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "with torch.no_grad():\n n_correct1 = 0\n n_correct2 = 0\n n_samples = 0\n for images, labels in test_loader:\n images = images.to(device)\n labels = labels.to(device)\n outputs1 = model1(images)\n outputs2 = model2(images)\n _, predictions1 = torch.max(outputs1, 1)\n _, predictions2 = torch.max(outputs2, 1)\n n_samples += labels.shape[0]\n n_correct1 += (predictions1 == labels).sum().item()\n n_correct2 += (predictions2 == labels).sum().item()\n acc1 = 100 * n_correct1 / n_samples\n acc2 = 100 * n_correct2 / n_samples" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Testing results (two separate models):\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "print(f'model1 accuracy = {acc1}%, model2 accuracy = {acc2}%')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Build the affinity matrix for graph matching\nAs shown in the following plot, the neural networks can be regarded as graphs. The weights corresponds to\nthe edge features, and the bias corresponds to the node features. In this example, the neural network\ndoes not have bias so that there are only edge features.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "plt.figure(figsize=(8, 4))\nimg = Image.open('../data/model_fusion.png')\nplt.imshow(img)\nplt.axis('off')\nst_time = time.perf_counter()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Define the graph matching affinity metric function\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "class Ground_Metric_GM:\n def __init__(self,\n model_1_param: torch.tensor = None,\n model_2_param: torch.tensor = None,\n conv_param: bool = False,\n bias_param: bool = False,\n pre_conv_param: bool = False,\n pre_conv_image_size_squared: int = None):\n self.model_1_param = model_1_param\n self.model_2_param = model_2_param\n self.conv_param = conv_param\n self.bias_param = bias_param\n # bias, or fully-connected from linear\n if bias_param is True or (conv_param is False and pre_conv_param is False):\n self.model_1_param = self.model_1_param.reshape(1, -1, 1)\n self.model_2_param = self.model_2_param.reshape(1, -1, 1)\n # fully-connected from conv\n elif conv_param is False and pre_conv_param is True:\n self.model_1_param = self.model_1_param.reshape(1, -1, pre_conv_image_size_squared)\n self.model_2_param = self.model_2_param.reshape(1, -1, pre_conv_image_size_squared)\n # conv\n else:\n self.model_1_param = self.model_1_param.reshape(1, -1, model_1_param.shape[-1])\n self.model_2_param = self.model_2_param.reshape(1, -1, model_2_param.shape[-1])\n\n def process_distance(self, p: int = 2):\n return torch.cdist(\n self.model_1_param.to(torch.float),\n self.model_2_param.to(torch.float),\n p=p)[0]\n\n def process_soft_affinity(self, p: int = 2):\n return torch.exp(0 - self.process_distance(p=p))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Define the affinity function between two neural networks. This function takes multiple neural network modules,\nand construct the corresponding affinity matrix which is further processed by the graph matching solver.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def graph_matching_fusion(networks: list):\n def total_node_num(network: torch.nn.Module):\n # count the total number of nodes in the network [network]\n num_nodes = 0\n for idx, (name, parameters) in enumerate(network.named_parameters()):\n if 'bias' in name:\n continue\n if idx == 0:\n num_nodes += parameters.shape[1]\n num_nodes += parameters.shape[0]\n return num_nodes\n\n n1 = total_node_num(network=networks[0])\n n2 = total_node_num(network=networks[1])\n assert (n1 == n2)\n affinity = torch.zeros([n1 * n2, n1 * n2], device=device)\n num_layers = len(list(zip(networks[0].parameters(), networks[1].parameters())))\n num_nodes_before = 0\n num_nodes_incremental = []\n num_nodes_layers = []\n pre_conv_list = []\n cur_conv_list = []\n conv_kernel_size_list = []\n num_nodes_pre = 0\n is_conv = False\n pre_conv = False\n pre_conv_out_channel = 1\n is_final_bias = False\n perm_is_complete = True\n named_weight_list_0 = [named_parameter for named_parameter in networks[0].named_parameters()]\n for idx, ((_, fc_layer0_weight), (_, fc_layer1_weight)) in \\\n enumerate(zip(networks[0].named_parameters(), networks[1].named_parameters())):\n assert fc_layer0_weight.shape == fc_layer1_weight.shape\n layer_shape = fc_layer0_weight.shape\n num_nodes_cur = fc_layer0_weight.shape[0]\n if len(layer_shape) > 1:\n if is_conv is True and len(layer_shape) == 2:\n num_nodes_pre = pre_conv_out_channel\n else:\n num_nodes_pre = fc_layer0_weight.shape[1]\n if idx >= 1 and len(named_weight_list_0[idx - 1][1].shape) == 1:\n pre_bias = True\n else:\n pre_bias = False\n if len(layer_shape) > 2:\n is_bias = False\n if not pre_bias:\n pre_conv = is_conv\n pre_conv_list.append(pre_conv)\n is_conv = True\n cur_conv_list.append(is_conv)\n fc_layer0_weight_data = fc_layer0_weight.data.view(\n fc_layer0_weight.shape[0], fc_layer0_weight.shape[1], -1)\n fc_layer1_weight_data = fc_layer1_weight.data.view(\n fc_layer1_weight.shape[0], fc_layer1_weight.shape[1], -1)\n elif len(layer_shape) == 2:\n is_bias = False\n if not pre_bias:\n pre_conv = is_conv\n pre_conv_list.append(pre_conv)\n is_conv = False\n cur_conv_list.append(is_conv)\n fc_layer0_weight_data = fc_layer0_weight.data\n fc_layer1_weight_data = fc_layer1_weight.data\n else:\n is_bias = True\n if not pre_bias:\n pre_conv = is_conv\n pre_conv_list.append(pre_conv)\n is_conv = False\n cur_conv_list.append(is_conv)\n fc_layer0_weight_data = fc_layer0_weight.data\n fc_layer1_weight_data = fc_layer1_weight.data\n if is_conv:\n pre_conv_out_channel = num_nodes_cur\n if is_bias is True and idx == num_layers - 1:\n is_final_bias = True\n if idx == 0:\n for a in range(num_nodes_pre):\n affinity[(num_nodes_before + a) * n2 + num_nodes_before + a] \\\n [(num_nodes_before + a) * n2 + num_nodes_before + a] \\\n = 1\n if idx == num_layers - 2 and 'bias' in named_weight_list_0[idx + 1][0] or \\\n idx == num_layers - 1 and 'bias' not in named_weight_list_0[idx][0]:\n for a in range(num_nodes_cur):\n affinity[(num_nodes_before + num_nodes_pre + a) * n2 + num_nodes_before + num_nodes_pre + a] \\\n [(num_nodes_before + num_nodes_pre + a) * n2 + num_nodes_before + num_nodes_pre + a] \\\n = 1\n if is_bias is False:\n ground_metric = Ground_Metric_GM(\n fc_layer0_weight_data, fc_layer1_weight_data, is_conv, is_bias,\n pre_conv, int(fc_layer0_weight_data.shape[1] / pre_conv_out_channel))\n else:\n ground_metric = Ground_Metric_GM(\n fc_layer0_weight_data, fc_layer1_weight_data, is_conv, is_bias,\n pre_conv, 1)\n\n layer_affinity = ground_metric.process_soft_affinity(p=2)\n\n if is_bias is False:\n pre_conv_kernel_size = fc_layer0_weight.shape[3] if is_conv else None\n conv_kernel_size_list.append(pre_conv_kernel_size)\n if is_bias is True and is_final_bias is False:\n for a in range(num_nodes_cur):\n for c in range(num_nodes_cur):\n affinity[(num_nodes_before + a) * n2 + num_nodes_before + c] \\\n [(num_nodes_before + a) * n2 + num_nodes_before + c] \\\n = layer_affinity[a][c]\n elif is_final_bias is False:\n for a in range(num_nodes_pre):\n for b in range(num_nodes_cur):\n affinity[\n (num_nodes_before + a) * n2 + num_nodes_before:\n (num_nodes_before + a) * n2 + num_nodes_before + num_nodes_pre,\n (num_nodes_before + num_nodes_pre + b) * n2 + num_nodes_before + num_nodes_pre:\n (num_nodes_before + num_nodes_pre + b) * n2 + num_nodes_before + num_nodes_pre + num_nodes_cur] \\\n = layer_affinity[a + b * num_nodes_pre].view(num_nodes_cur, num_nodes_pre).transpose(0, 1)\n if is_bias is False:\n num_nodes_before += num_nodes_pre\n num_nodes_incremental.append(num_nodes_before)\n num_nodes_layers.append(num_nodes_cur)\n # affinity = (affinity + affinity.t()) / 2\n return affinity, [n1, n2, num_nodes_incremental, num_nodes_layers, cur_conv_list, conv_kernel_size_list]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Get the affinity (similarity) matrix between model1 and model2.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "K, params = graph_matching_fusion([model1, model2])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Align the models by graph matching\nAlign the channels of model1 & model2 by maximize the affinity (similarity) via graph matching algorithms.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "n1 = params[0]\nn2 = params[1]\nX = pygm.sm(K, n1, n2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Project ``X`` to neural network matching result. The neural network matching matrix is built by applying\nHungarian to small blocks of ``X``, because only the channels from the same neural network layer can be\nmatched.\n\n

Note

In this example, we assume the last FC layer is aligned and need not to be matched.

\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "new_X = torch.zeros_like(X)\nnew_X[:params[2][0], :params[2][0]] = torch.eye(params[2][0], device=device)\nfor start_idx, length in zip(params[2][:-1], params[3][:-1]): # params[2] and params[3] are the indices of layers\n slicing = slice(start_idx, start_idx + length)\n new_X[slicing, slicing] = pygm.hungarian(X[slicing, slicing])\n# assume the last FC layer is aligned\nslicing = slice(params[2][-1], params[2][-1] + params[3][-1])\nnew_X[slicing, slicing] = torch.eye(params[3][-1], device=device)\nX = new_X" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Visualization of the matching result. The black lines splits the channels of different layers.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "plt.figure(figsize=(4, 4))\nplt.imshow(X.cpu().numpy(), cmap='Blues')\nfor idx in params[2]:\n plt.axvline(x=idx, color='k')\n plt.axhline(y=idx, color='k')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Define the alignment function: fuse the models based on matching result\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def align(solution, fusion_proportion, networks: list, params: list):\n [_, _, num_nodes_incremental, num_nodes_layers, cur_conv_list, conv_kernel_size_list] = params\n named_weight_list_0 = [named_parameter for named_parameter in networks[0].named_parameters()]\n aligned_wt_0 = [parameter.data for name, parameter in named_weight_list_0]\n idx = 0\n num_layers = len(aligned_wt_0)\n for num_before, num_cur, cur_conv, cur_kernel_size in \\\n zip(num_nodes_incremental, num_nodes_layers, cur_conv_list, conv_kernel_size_list):\n perm = solution[num_before:num_before + num_cur, num_before:num_before + num_cur]\n assert 'bias' not in named_weight_list_0[idx][0]\n if len(named_weight_list_0[idx][1].shape) == 4:\n aligned_wt_0[idx] = (perm.transpose(0, 1).to(torch.float64) @\n aligned_wt_0[idx].to(torch.float64).permute(2, 3, 0, 1)) \\\n .permute(2, 3, 0, 1)\n else:\n aligned_wt_0[idx] = perm.transpose(0, 1).to(torch.float64) @ aligned_wt_0[idx].to(torch.float64)\n idx += 1\n if idx >= num_layers:\n continue\n if 'bias' in named_weight_list_0[idx][0]:\n aligned_wt_0[idx] = aligned_wt_0[idx].to(torch.float64) @ perm.to(torch.float64)\n idx += 1\n if idx >= num_layers:\n continue\n if cur_conv and len(named_weight_list_0[idx][1].shape) == 2:\n aligned_wt_0[idx] = (aligned_wt_0[idx].to(torch.float64)\n .reshape(aligned_wt_0[idx].shape[0], 64, -1)\n .permute(0, 2, 1)\n @ perm.to(torch.float64)) \\\n .permute(0, 2, 1) \\\n .reshape(aligned_wt_0[idx].shape[0], -1)\n elif len(named_weight_list_0[idx][1].shape) == 4:\n aligned_wt_0[idx] = (aligned_wt_0[idx].to(torch.float64)\n .permute(2, 3, 0, 1)\n @ perm.to(torch.float64)) \\\n .permute(2, 3, 0, 1)\n else:\n aligned_wt_0[idx] = aligned_wt_0[idx].to(torch.float64) @ perm.to(torch.float64)\n assert idx == num_layers\n\n averaged_weights = []\n for idx, parameter in enumerate(networks[1].parameters()):\n averaged_weights.append((1 - fusion_proportion) * aligned_wt_0[idx] + fusion_proportion * parameter)\n return averaged_weights" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Test the fused model\nThe ``fusion_proportion`` variable denotes the contribution to the new model. For example, if ``fusion_proportion=0.2``,\nthe fused model = 80% model1 + 20% model2.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def align_model_and_test(X):\n acc_list = []\n for fusion_proportion in torch.arange(0, 1.1, 0.1):\n fused_weights = align(X, fusion_proportion, [model1, model2], params)\n\n fused_model = SimpleNet()\n state_dict = fused_model.state_dict()\n for idx, (key, _) in enumerate(state_dict.items()):\n state_dict[key] = fused_weights[idx]\n fused_model.load_state_dict(state_dict)\n fused_model.to(device)\n test_loss = 0\n correct = 0\n for data, target in test_loader:\n data = data.to(device)\n target = target.to(device)\n output = fused_model(data)\n test_loss += F.nll_loss(output, target, reduction='sum').item()\n pred = output.data.max(1, keepdim=True)[1]\n correct += pred.eq(target.data.view_as(pred)).sum()\n test_loss /= len(test_loader.dataset)\n acc = 100. * correct / len(test_loader.dataset)\n print(\n f\"{1 - fusion_proportion:.2f} model1 + {fusion_proportion:.2f} model2 -> fused model accuracy: {acc:.2f}%\")\n acc_list.append(acc)\n return torch.tensor(acc_list)\n\n\nprint('Graph Matching Fusion')\ngm_acc_list = align_model_and_test(X)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Compare with vanilla model fusion (no matching), graph matching method stabilizes the fusion step:\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "print('No Matching Fusion')\nvanilla_acc_list = align_model_and_test(torch.eye(n1, device=device))\n\nplt.figure(figsize=(4, 4))\nplt.title('Fused Model Accuracy')\nplt.plot(torch.arange(0, 1.1, 0.1).numpy(), gm_acc_list.cpu().numpy(), 'r*-', label='Graph Matching Fusion')\nplt.plot(torch.arange(0, 1.1, 0.1).numpy(), vanilla_acc_list.cpu().numpy(), 'b*-', label='No Matching Fusion')\nplt.plot(torch.arange(0, 1.1, 0.1).numpy(), [acc1] * 11, '--', color=\"gray\", label='Model1 Accuracy')\nplt.plot(torch.arange(0, 1.1, 0.1).numpy(), [acc2] * 11, '--', color=\"brown\", label='Model2 Accuracy')\nplt.gca().set_xlabel('Fusion Proportion')\nplt.gca().set_ylabel('Accuracy (%)')\nplt.ylim((70, 87))\nplt.legend(loc=3)\nplt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Print the result summary\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "end_time = time.perf_counter()\nprint(f'time consumed for model fusion: {end_time - st_time:.2f} seconds')\nprint(f'model1 accuracy = {acc1}%, model2 accuracy = {acc2}%')\nprint(f\"best fused model accuracy: {torch.max(gm_acc_list):.2f}%\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "

Note

This example supports both GPU and CPU, and the online documentation is built by a CPU-only machine.\n The efficiency will be significantly improved if you run this code on GPU.

\n\n\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "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.8.0" } }, "nbformat": 4, "nbformat_minor": 0 }