mynote.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "from torchvision.models.resnet import resnet50, Bottleneck\n",
    "import torch.fx as fx\n",
    "\n",
    "import torch\n",
    "import torch._dynamo\n",
    "import numpy as np\n",
    "import pickle\n",
    "from math import ceil\n",
    "from typing import List"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "0: Load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unpickle(file):\n",
    "    with open(file, 'rb') as fo:\n",
    "        dict = pickle.load(fo, encoding='latin1')\n",
    "    return dict\n",
    "\n",
    "data_batch_1 = unpickle('/home/lyy/data/cifar-10-batches-py/data_batch_1')\n",
    "data_batch_2 = unpickle('/home/lyy/data/cifar-10-batches-py/data_batch_2')\n",
    "data_batch_3 = unpickle('/home/lyy/data/cifar-10-batches-py/data_batch_3')\n",
    "data_batch_4 = unpickle('/home/lyy/data/cifar-10-batches-py/data_batch_4')\n",
    "data_batch_5 = unpickle('/home/lyy/data/cifar-10-batches-py/data_batch_5')\n",
    "test_batch = unpickle('/home/lyy/data/cifar-10-batches-py/test_batch')\n",
    "\n",
    "\n",
    "X_train_data = np.concatenate((data_batch_1['data'], data_batch_2['data'], data_batch_3['data'], data_batch_4['data'], data_batch_5['data']))\n",
    "X_train_labels = data_batch_1['labels'] + data_batch_2['labels'] + data_batch_3['labels'] + data_batch_4['labels'] + data_batch_5['labels']\n",
    "X_test_data = test_batch['data']\n",
    "X_test_labels = test_batch['labels']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1: preprocess the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train_data = X_train_data.reshape(len(X_train_data),3,32,32)\n",
    "\n",
    "batch_size = 128\n",
    "split_data_list = np.array_split(X_train_data,ceil(X_train_data.shape[0]/batch_size), axis=0)\n",
    "how_many_batches = len(split_data_list)\n",
    "split_labels_list = np.array_split(X_train_labels,ceil(len(X_train_labels)/batch_size), axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2 import model and define optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/lyy/.local/lib/python3.8/site-packages/torchvision/models/_utils.py:135: UserWarning: Using 'weights' as positional parameter(s) is deprecated since 0.13 and may be removed in the future. Please use keyword parameter(s) instead.\n",
      "  warnings.warn(\n",
      "/home/lyy/.local/lib/python3.8/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=ResNet50_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet50_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "net = resnet50(10).to('cuda')\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.1, momentum=0.9, weight_decay=0.0001)\n",
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor = 0.1, patience=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "    # tranform data type\n",
    "    # inputs = torch.from_numpy(inputs)\n",
    "    # inputs = inputs.float()\n",
    "    # inputs = inputs.to('cuda')\n",
    "\n",
    "    # labels = torch.from_numpy(labels)\n",
    "    # labels = labels.to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "def segment(inputs, labels, losses):\n",
    "    # tranform data type\n",
    "    inputs = torch.from_numpy(inputs)\n",
    "    inputs = inputs.float()\n",
    "    inputs = inputs.to('cuda')\n",
    "\n",
    "    labels = torch.from_numpy(labels)\n",
    "    labels = labels.to('cuda')\n",
    "    optimizer.zero_grad()\n",
    "    outputs = net(inputs)\n",
    "    loss = criterion(outputs, labels)\n",
    "    losses.append(loss.item())\n",
    "\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4: start to train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 1 # actual:200\n",
    "for epoch in range(EPOCHS):\n",
    "    losses = []\n",
    "    running_loss = 0\n",
    "\n",
    "    for i in range(how_many_batches):\n",
    "        # move to the segment funciton\n",
    "        loss = segment(split_data_list[i], split_labels_list[i], losses)\n",
    "        # prove of concept: can be put into the loop and print everything\n",
    "        # print(opt_model_segment(split_data_list[i], split_labels_list[i], losses))\n",
    "\n",
    "        running_loss += loss\n",
    "        \n",
    "        if i%100 == 0 and i > 0:\n",
    "            print(f'Loss [{epoch+1}, {i}](epoch, minibatch): ', running_loss / 100)\n",
    "            running_loss = 0.0\n",
    "\n",
    "\n",
    "    avg_loss = sum(losses)/len(losses)\n",
    "    scheduler.step(avg_loss)\n",
    "\n",
    "print('Training Done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "def custom_backend(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):\n",
    "    print(\"custom backend called wiht FX graph:\")\n",
    "\n",
    "    gm.graph.print_tabular()\n",
    "    return gm.forward\n",
    "\n",
    "torch._dynamo.reset()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6: print out the graph of the torch.nn.module(aka resnet50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opcode         name                   target                                                      args                                   kwargs\n",
      "-------------  ---------------------  ----------------------------------------------------------  -------------------------------------  --------\n",
      "placeholder    x                      x                                                           ()                                     {}\n",
      "call_module    conv1                  conv1                                                       (x,)                                   {}\n",
      "call_module    bn1                    bn1                                                         (conv1,)                               {}\n",
      "call_module    relu                   relu                                                        (bn1,)                                 {}\n",
      "call_module    maxpool                maxpool                                                     (relu,)                                {}\n",
      "call_module    layer1_0_conv1         layer1.0.conv1                                              (maxpool,)                             {}\n",
      "call_module    layer1_0_bn1           layer1.0.bn1                                                (layer1_0_conv1,)                      {}\n",
      "call_module    layer1_0_relu          layer1.0.relu                                               (layer1_0_bn1,)                        {}\n",
      "call_module    layer1_0_conv2         layer1.0.conv2                                              (layer1_0_relu,)                       {}\n",
      "call_module    layer1_0_bn2           layer1.0.bn2                                                (layer1_0_conv2,)                      {}\n",
      "call_module    layer1_0_relu_1        layer1.0.relu                                               (layer1_0_bn2,)                        {}\n",
      "call_module    layer1_0_conv3         layer1.0.conv3                                              (layer1_0_relu_1,)                     {}\n",
      "call_module    layer1_0_bn3           layer1.0.bn3                                                (layer1_0_conv3,)                      {}\n",
      "call_module    layer1_0_downsample_0  layer1.0.downsample.0                                       (maxpool,)                             {}\n",
      "call_module    layer1_0_downsample_1  layer1.0.downsample.1                                       (layer1_0_downsample_0,)               {}\n",
      "call_function  add                    <built-in function add>                                     (layer1_0_bn3, layer1_0_downsample_1)  {}\n",
      "call_module    layer1_0_relu_2        layer1.0.relu                                               (add,)                                 {}\n",
      "call_module    layer1_1_conv1         layer1.1.conv1                                              (layer1_0_relu_2,)                     {}\n",
      "call_module    layer1_1_bn1           layer1.1.bn1                                                (layer1_1_conv1,)                      {}\n",
      "call_module    layer1_1_relu          layer1.1.relu                                               (layer1_1_bn1,)                        {}\n",
      "call_module    layer1_1_conv2         layer1.1.conv2                                              (layer1_1_relu,)                       {}\n",
      "call_module    layer1_1_bn2           layer1.1.bn2                                                (layer1_1_conv2,)                      {}\n",
      "call_module    layer1_1_relu_1        layer1.1.relu                                               (layer1_1_bn2,)                        {}\n",
      "call_module    layer1_1_conv3         layer1.1.conv3                                              (layer1_1_relu_1,)                     {}\n",
      "call_module    layer1_1_bn3           layer1.1.bn3                                                (layer1_1_conv3,)                      {}\n",
      "call_function  add_1                  <built-in function add>                                     (layer1_1_bn3, layer1_0_relu_2)        {}\n",
      "call_module    layer1_1_relu_2        layer1.1.relu                                               (add_1,)                               {}\n",
      "call_module    layer1_2_conv1         layer1.2.conv1                                              (layer1_1_relu_2,)                     {}\n",
      "call_module    layer1_2_bn1           layer1.2.bn1                                                (layer1_2_conv1,)                      {}\n",
      "call_module    layer1_2_relu          layer1.2.relu                                               (layer1_2_bn1,)                        {}\n",
      "call_module    layer1_2_conv2         layer1.2.conv2                                              (layer1_2_relu,)                       {}\n",
      "call_module    layer1_2_bn2           layer1.2.bn2                                                (layer1_2_conv2,)                      {}\n",
      "call_module    layer1_2_relu_1        layer1.2.relu                                               (layer1_2_bn2,)                        {}\n",
      "call_module    layer1_2_conv3         layer1.2.conv3                                              (layer1_2_relu_1,)                     {}\n",
      "call_module    layer1_2_bn3           layer1.2.bn3                                                (layer1_2_conv3,)                      {}\n",
      "call_function  add_2                  <built-in function add>                                     (layer1_2_bn3, layer1_1_relu_2)        {}\n",
      "call_module    layer1_2_relu_2        layer1.2.relu                                               (add_2,)                               {}\n",
      "call_module    layer2_0_conv1         layer2.0.conv1                                              (layer1_2_relu_2,)                     {}\n",
      "call_module    layer2_0_bn1           layer2.0.bn1                                                (layer2_0_conv1,)                      {}\n",
      "call_module    layer2_0_relu          layer2.0.relu                                               (layer2_0_bn1,)                        {}\n",
      "call_module    layer2_0_conv2         layer2.0.conv2                                              (layer2_0_relu,)                       {}\n",
      "call_module    layer2_0_bn2           layer2.0.bn2                                                (layer2_0_conv2,)                      {}\n",
      "call_module    layer2_0_relu_1        layer2.0.relu                                               (layer2_0_bn2,)                        {}\n",
      "call_module    layer2_0_conv3         layer2.0.conv3                                              (layer2_0_relu_1,)                     {}\n",
      "call_module    layer2_0_bn3           layer2.0.bn3                                                (layer2_0_conv3,)                      {}\n",
      "call_module    layer2_0_downsample_0  layer2.0.downsample.0                                       (layer1_2_relu_2,)                     {}\n",
      "call_module    layer2_0_downsample_1  layer2.0.downsample.1                                       (layer2_0_downsample_0,)               {}\n",
      "call_function  add_3                  <built-in function add>                                     (layer2_0_bn3, layer2_0_downsample_1)  {}\n",
      "call_module    layer2_0_relu_2        layer2.0.relu                                               (add_3,)                               {}\n",
      "call_module    layer2_1_conv1         layer2.1.conv1                                              (layer2_0_relu_2,)                     {}\n",
      "call_module    layer2_1_bn1           layer2.1.bn1                                                (layer2_1_conv1,)                      {}\n",
      "call_module    layer2_1_relu          layer2.1.relu                                               (layer2_1_bn1,)                        {}\n",
      "call_module    layer2_1_conv2         layer2.1.conv2                                              (layer2_1_relu,)                       {}\n",
      "call_module    layer2_1_bn2           layer2.1.bn2                                                (layer2_1_conv2,)                      {}\n",
      "call_module    layer2_1_relu_1        layer2.1.relu                                               (layer2_1_bn2,)                        {}\n",
      "call_module    layer2_1_conv3         layer2.1.conv3                                              (layer2_1_relu_1,)                     {}\n",
      "call_module    layer2_1_bn3           layer2.1.bn3                                                (layer2_1_conv3,)                      {}\n",
      "call_function  add_4                  <built-in function add>                                     (layer2_1_bn3, layer2_0_relu_2)        {}\n",
      "call_module    layer2_1_relu_2        layer2.1.relu                                               (add_4,)                               {}\n",
      "call_module    layer2_2_conv1         layer2.2.conv1                                              (layer2_1_relu_2,)                     {}\n",
      "call_module    layer2_2_bn1           layer2.2.bn1                                                (layer2_2_conv1,)                      {}\n",
      "call_module    layer2_2_relu          layer2.2.relu                                               (layer2_2_bn1,)                        {}\n",
      "call_module    layer2_2_conv2         layer2.2.conv2                                              (layer2_2_relu,)                       {}\n",
      "call_module    layer2_2_bn2           layer2.2.bn2                                                (layer2_2_conv2,)                      {}\n",
      "call_module    layer2_2_relu_1        layer2.2.relu                                               (layer2_2_bn2,)                        {}\n",
      "call_module    layer2_2_conv3         layer2.2.conv3                                              (layer2_2_relu_1,)                     {}\n",
      "call_module    layer2_2_bn3           layer2.2.bn3                                                (layer2_2_conv3,)                      {}\n",
      "call_function  add_5                  <built-in function add>                                     (layer2_2_bn3, layer2_1_relu_2)        {}\n",
      "call_module    layer2_2_relu_2        layer2.2.relu                                               (add_5,)                               {}\n",
      "call_module    layer2_3_conv1         layer2.3.conv1                                              (layer2_2_relu_2,)                     {}\n",
      "call_module    layer2_3_bn1           layer2.3.bn1                                                (layer2_3_conv1,)                      {}\n",
      "call_module    layer2_3_relu          layer2.3.relu                                               (layer2_3_bn1,)                        {}\n",
      "call_module    layer2_3_conv2         layer2.3.conv2                                              (layer2_3_relu,)                       {}\n",
      "call_module    layer2_3_bn2           layer2.3.bn2                                                (layer2_3_conv2,)                      {}\n",
      "call_module    layer2_3_relu_1        layer2.3.relu                                               (layer2_3_bn2,)                        {}\n",
      "call_module    layer2_3_conv3         layer2.3.conv3                                              (layer2_3_relu_1,)                     {}\n",
      "call_module    layer2_3_bn3           layer2.3.bn3                                                (layer2_3_conv3,)                      {}\n",
      "call_function  add_6                  <built-in function add>                                     (layer2_3_bn3, layer2_2_relu_2)        {}\n",
      "call_module    layer2_3_relu_2        layer2.3.relu                                               (add_6,)                               {}\n",
      "call_module    layer3_0_conv1         layer3.0.conv1                                              (layer2_3_relu_2,)                     {}\n",
      "call_module    layer3_0_bn1           layer3.0.bn1                                                (layer3_0_conv1,)                      {}\n",
      "call_module    layer3_0_relu          layer3.0.relu                                               (layer3_0_bn1,)                        {}\n",
      "call_module    layer3_0_conv2         layer3.0.conv2                                              (layer3_0_relu,)                       {}\n",
      "call_module    layer3_0_bn2           layer3.0.bn2                                                (layer3_0_conv2,)                      {}\n",
      "call_module    layer3_0_relu_1        layer3.0.relu                                               (layer3_0_bn2,)                        {}\n",
      "call_module    layer3_0_conv3         layer3.0.conv3                                              (layer3_0_relu_1,)                     {}\n",
      "call_module    layer3_0_bn3           layer3.0.bn3                                                (layer3_0_conv3,)                      {}\n",
      "call_module    layer3_0_downsample_0  layer3.0.downsample.0                                       (layer2_3_relu_2,)                     {}\n",
      "call_module    layer3_0_downsample_1  layer3.0.downsample.1                                       (layer3_0_downsample_0,)               {}\n",
      "call_function  add_7                  <built-in function add>                                     (layer3_0_bn3, layer3_0_downsample_1)  {}\n",
      "call_module    layer3_0_relu_2        layer3.0.relu                                               (add_7,)                               {}\n",
      "call_module    layer3_1_conv1         layer3.1.conv1                                              (layer3_0_relu_2,)                     {}\n",
      "call_module    layer3_1_bn1           layer3.1.bn1                                                (layer3_1_conv1,)                      {}\n",
      "call_module    layer3_1_relu          layer3.1.relu                                               (layer3_1_bn1,)                        {}\n",
      "call_module    layer3_1_conv2         layer3.1.conv2                                              (layer3_1_relu,)                       {}\n",
      "call_module    layer3_1_bn2           layer3.1.bn2                                                (layer3_1_conv2,)                      {}\n",
      "call_module    layer3_1_relu_1        layer3.1.relu                                               (layer3_1_bn2,)                        {}\n",
      "call_module    layer3_1_conv3         layer3.1.conv3                                              (layer3_1_relu_1,)                     {}\n",
      "call_module    layer3_1_bn3           layer3.1.bn3                                                (layer3_1_conv3,)                      {}\n",
      "call_function  add_8                  <built-in function add>                                     (layer3_1_bn3, layer3_0_relu_2)        {}\n",
      "call_module    layer3_1_relu_2        layer3.1.relu                                               (add_8,)                               {}\n",
      "call_module    layer3_2_conv1         layer3.2.conv1                                              (layer3_1_relu_2,)                     {}\n",
      "call_module    layer3_2_bn1           layer3.2.bn1                                                (layer3_2_conv1,)                      {}\n",
      "call_module    layer3_2_relu          layer3.2.relu                                               (layer3_2_bn1,)                        {}\n",
      "call_module    layer3_2_conv2         layer3.2.conv2                                              (layer3_2_relu,)                       {}\n",
      "call_module    layer3_2_bn2           layer3.2.bn2                                                (layer3_2_conv2,)                      {}\n",
      "call_module    layer3_2_relu_1        layer3.2.relu                                               (layer3_2_bn2,)                        {}\n",
      "call_module    layer3_2_conv3         layer3.2.conv3                                              (layer3_2_relu_1,)                     {}\n",
      "call_module    layer3_2_bn3           layer3.2.bn3                                                (layer3_2_conv3,)                      {}\n",
      "call_function  add_9                  <built-in function add>                                     (layer3_2_bn3, layer3_1_relu_2)        {}\n",
      "call_module    layer3_2_relu_2        layer3.2.relu                                               (add_9,)                               {}\n",
      "call_module    layer3_3_conv1         layer3.3.conv1                                              (layer3_2_relu_2,)                     {}\n",
      "call_module    layer3_3_bn1           layer3.3.bn1                                                (layer3_3_conv1,)                      {}\n",
      "call_module    layer3_3_relu          layer3.3.relu                                               (layer3_3_bn1,)                        {}\n",
      "call_module    layer3_3_conv2         layer3.3.conv2                                              (layer3_3_relu,)                       {}\n",
      "call_module    layer3_3_bn2           layer3.3.bn2                                                (layer3_3_conv2,)                      {}\n",
      "call_module    layer3_3_relu_1        layer3.3.relu                                               (layer3_3_bn2,)                        {}\n",
      "call_module    layer3_3_conv3         layer3.3.conv3                                              (layer3_3_relu_1,)                     {}\n",
      "call_module    layer3_3_bn3           layer3.3.bn3                                                (layer3_3_conv3,)                      {}\n",
      "call_function  add_10                 <built-in function add>                                     (layer3_3_bn3, layer3_2_relu_2)        {}\n",
      "call_module    layer3_3_relu_2        layer3.3.relu                                               (add_10,)                              {}\n",
      "call_module    layer3_4_conv1         layer3.4.conv1                                              (layer3_3_relu_2,)                     {}\n",
      "call_module    layer3_4_bn1           layer3.4.bn1                                                (layer3_4_conv1,)                      {}\n",
      "call_module    layer3_4_relu          layer3.4.relu                                               (layer3_4_bn1,)                        {}\n",
      "call_module    layer3_4_conv2         layer3.4.conv2                                              (layer3_4_relu,)                       {}\n",
      "call_module    layer3_4_bn2           layer3.4.bn2                                                (layer3_4_conv2,)                      {}\n",
      "call_module    layer3_4_relu_1        layer3.4.relu                                               (layer3_4_bn2,)                        {}\n",
      "call_module    layer3_4_conv3         layer3.4.conv3                                              (layer3_4_relu_1,)                     {}\n",
      "call_module    layer3_4_bn3           layer3.4.bn3                                                (layer3_4_conv3,)                      {}\n",
      "call_function  add_11                 <built-in function add>                                     (layer3_4_bn3, layer3_3_relu_2)        {}\n",
      "call_module    layer3_4_relu_2        layer3.4.relu                                               (add_11,)                              {}\n",
      "call_module    layer3_5_conv1         layer3.5.conv1                                              (layer3_4_relu_2,)                     {}\n",
      "call_module    layer3_5_bn1           layer3.5.bn1                                                (layer3_5_conv1,)                      {}\n",
      "call_module    layer3_5_relu          layer3.5.relu                                               (layer3_5_bn1,)                        {}\n",
      "call_module    layer3_5_conv2         layer3.5.conv2                                              (layer3_5_relu,)                       {}\n",
      "call_module    layer3_5_bn2           layer3.5.bn2                                                (layer3_5_conv2,)                      {}\n",
      "call_module    layer3_5_relu_1        layer3.5.relu                                               (layer3_5_bn2,)                        {}\n",
      "call_module    layer3_5_conv3         layer3.5.conv3                                              (layer3_5_relu_1,)                     {}\n",
      "call_module    layer3_5_bn3           layer3.5.bn3                                                (layer3_5_conv3,)                      {}\n",
      "call_function  add_12                 <built-in function add>                                     (layer3_5_bn3, layer3_4_relu_2)        {}\n",
      "call_module    layer3_5_relu_2        layer3.5.relu                                               (add_12,)                              {}\n",
      "call_module    layer4_0_conv1         layer4.0.conv1                                              (layer3_5_relu_2,)                     {}\n",
      "call_module    layer4_0_bn1           layer4.0.bn1                                                (layer4_0_conv1,)                      {}\n",
      "call_module    layer4_0_relu          layer4.0.relu                                               (layer4_0_bn1,)                        {}\n",
      "call_module    layer4_0_conv2         layer4.0.conv2                                              (layer4_0_relu,)                       {}\n",
      "call_module    layer4_0_bn2           layer4.0.bn2                                                (layer4_0_conv2,)                      {}\n",
      "call_module    layer4_0_relu_1        layer4.0.relu                                               (layer4_0_bn2,)                        {}\n",
      "call_module    layer4_0_conv3         layer4.0.conv3                                              (layer4_0_relu_1,)                     {}\n",
      "call_module    layer4_0_bn3           layer4.0.bn3                                                (layer4_0_conv3,)                      {}\n",
      "call_module    layer4_0_downsample_0  layer4.0.downsample.0                                       (layer3_5_relu_2,)                     {}\n",
      "call_module    layer4_0_downsample_1  layer4.0.downsample.1                                       (layer4_0_downsample_0,)               {}\n",
      "call_function  add_13                 <built-in function add>                                     (layer4_0_bn3, layer4_0_downsample_1)  {}\n",
      "call_module    layer4_0_relu_2        layer4.0.relu                                               (add_13,)                              {}\n",
      "call_module    layer4_1_conv1         layer4.1.conv1                                              (layer4_0_relu_2,)                     {}\n",
      "call_module    layer4_1_bn1           layer4.1.bn1                                                (layer4_1_conv1,)                      {}\n",
      "call_module    layer4_1_relu          layer4.1.relu                                               (layer4_1_bn1,)                        {}\n",
      "call_module    layer4_1_conv2         layer4.1.conv2                                              (layer4_1_relu,)                       {}\n",
      "call_module    layer4_1_bn2           layer4.1.bn2                                                (layer4_1_conv2,)                      {}\n",
      "call_module    layer4_1_relu_1        layer4.1.relu                                               (layer4_1_bn2,)                        {}\n",
      "call_module    layer4_1_conv3         layer4.1.conv3                                              (layer4_1_relu_1,)                     {}\n",
      "call_module    layer4_1_bn3           layer4.1.bn3                                                (layer4_1_conv3,)                      {}\n",
      "call_function  add_14                 <built-in function add>                                     (layer4_1_bn3, layer4_0_relu_2)        {}\n",
      "call_module    layer4_1_relu_2        layer4.1.relu                                               (add_14,)                              {}\n",
      "call_module    layer4_2_conv1         layer4.2.conv1                                              (layer4_1_relu_2,)                     {}\n",
      "call_module    layer4_2_bn1           layer4.2.bn1                                                (layer4_2_conv1,)                      {}\n",
      "call_module    layer4_2_relu          layer4.2.relu                                               (layer4_2_bn1,)                        {}\n",
      "call_module    layer4_2_conv2         layer4.2.conv2                                              (layer4_2_relu,)                       {}\n",
      "call_module    layer4_2_bn2           layer4.2.bn2                                                (layer4_2_conv2,)                      {}\n",
      "call_module    layer4_2_relu_1        layer4.2.relu                                               (layer4_2_bn2,)                        {}\n",
      "call_module    layer4_2_conv3         layer4.2.conv3                                              (layer4_2_relu_1,)                     {}\n",
      "call_module    layer4_2_bn3           layer4.2.bn3                                                (layer4_2_conv3,)                      {}\n",
      "call_function  add_15                 <built-in function add>                                     (layer4_2_bn3, layer4_1_relu_2)        {}\n",
      "call_module    layer4_2_relu_2        layer4.2.relu                                               (add_15,)                              {}\n",
      "call_module    avgpool                avgpool                                                     (layer4_2_relu_2,)                     {}\n",
      "call_function  flatten                <built-in method flatten of type object at 0x7ff353e9d100>  (avgpool, 1)                           {}\n",
      "call_module    fc                     fc                                                          (flatten,)                             {}\n",
      "output         output                 output                                                      (fc,)                                  {}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/lyy/.local/lib/python3.8/site-packages/torch/overrides.py:111: UserWarning: 'has_cuda' is deprecated, please use 'torch.backends.cuda.is_built()'\n",
      "  torch.has_cuda,\n",
      "/home/lyy/.local/lib/python3.8/site-packages/torch/overrides.py:112: UserWarning: 'has_cudnn' is deprecated, please use 'torch.backends.cudnn.is_available()'\n",
      "  torch.has_cudnn,\n",
      "/home/lyy/.local/lib/python3.8/site-packages/torch/overrides.py:118: UserWarning: 'has_mps' is deprecated, please use 'torch.backends.mps.is_built()'\n",
      "  torch.has_mps,\n",
      "/home/lyy/.local/lib/python3.8/site-packages/torch/overrides.py:119: UserWarning: 'has_mkldnn' is deprecated, please use 'torch.backends.mkldnn.is_available()'\n",
      "  torch.has_mkldnn,\n"
     ]
    }
   ],
   "source": [
    "#gm = torch.fx.symbolic_trace(m)\n",
    "gm = torch.fx.symbolic_trace(net)\n",
    "# call and print the graph\n",
    "gm.graph.print_tabular()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: wouldn't work, do not know why yet\n",
    "opt_model = torch.compile(net, backend=custom_backend)\n",
    "print(opt_model(split_data_list[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "7: print out the graph of segment(only 1 slice of data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "size of the split_data_list[0] is (128, 3, 32, 32) which is splitted with mini batch size = 128\n",
    "size of the split_label_list[0] is (128,), representing the data point of the 128 data points\n",
    "losses is initialized to an empty list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(128,)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "split_labels_list[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.Tensor'>\n"
     ]
    }
   ],
   "source": [
    "print(type(inputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "custom backend called wiht FX graph:\n",
      "opcode       name       target     args               kwargs\n",
      "-----------  ---------  ---------  -----------------  --------\n",
      "placeholder  l_stack0_  L_stack0_  ()                 {}\n",
      "call_method  float_1    float      (l_stack0_,)       {}\n",
      "call_method  to         to         (float_1, 'cuda')  {}\n",
      "output       output     output     ((to,),)           {}\n",
      "custom backend called wiht FX graph:\n",
      "opcode       name       target     args                 kwargs\n",
      "-----------  ---------  ---------  -------------------  --------\n",
      "placeholder  l_stack0_  L_stack0_  ()                   {}\n",
      "call_method  to         to         (l_stack0_, 'cuda')  {}\n",
      "output       output     output     ((to,),)             {}\n",
      "custom backend called wiht FX graph:\n",
      "opcode         name                                        target                                                      args                                                                             kwargs\n",
      "-------------  ------------------------------------------  ----------------------------------------------------------  -------------------------------------------------------------------------------  --------\n",
      "placeholder    l_inputs_                                   L_inputs_                                                   ()                                                                               {}\n",
      "placeholder    l_labels_                                   L_labels_                                                   ()                                                                               {}\n",
      "call_module    g__net___conv1                              G__net___conv1                                              (l_inputs_,)                                                                     {}\n",
      "call_module    g__net___bn1                                G__net___bn1                                                (g__net___conv1,)                                                                {}\n",
      "call_module    g__net___relu                               G__net___relu                                               (g__net___bn1,)                                                                  {}\n",
      "call_module    g__net___maxpool                            G__net___maxpool                                            (g__net___relu,)                                                                 {}\n",
      "call_module    getattr_g__net___layer1___0___conv1         getattr_G__net___layer1___0___conv1                         (g__net___maxpool,)                                                              {}\n",
      "call_module    getattr_g__net___layer1___0___bn1           getattr_G__net___layer1___0___bn1                           (getattr_g__net___layer1___0___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer1___0___relu          getattr_G__net___layer1___0___relu                          (getattr_g__net___layer1___0___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer1___0___conv2         getattr_G__net___layer1___0___conv2                         (getattr_g__net___layer1___0___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer1___0___bn2           getattr_G__net___layer1___0___bn2                           (getattr_g__net___layer1___0___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer1___0___relu_1        getattr_G__net___layer1___0___relu                          (getattr_g__net___layer1___0___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer1___0___conv3         getattr_G__net___layer1___0___conv3                         (getattr_g__net___layer1___0___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer1___0___bn3           getattr_G__net___layer1___0___bn3                           (getattr_g__net___layer1___0___conv3,)                                           {}\n",
      "call_module    getattr_g__net___layer1___0___downsample_0  getattr_G__net___layer1___0___downsample_0                  (g__net___maxpool,)                                                              {}\n",
      "call_module    getattr_g__net___layer1___0___downsample_1  getattr_G__net___layer1___0___downsample_1                  (getattr_g__net___layer1___0___downsample_0,)                                    {}\n",
      "call_function  iadd                                        <built-in function iadd>                                    (getattr_g__net___layer1___0___bn3, getattr_g__net___layer1___0___downsample_1)  {}\n",
      "call_module    getattr_g__net___layer1___0___relu_2        getattr_G__net___layer1___0___relu                          (iadd,)                                                                          {}\n",
      "call_module    getattr_g__net___layer1___1___conv1         getattr_G__net___layer1___1___conv1                         (getattr_g__net___layer1___0___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer1___1___bn1           getattr_G__net___layer1___1___bn1                           (getattr_g__net___layer1___1___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer1___1___relu          getattr_G__net___layer1___1___relu                          (getattr_g__net___layer1___1___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer1___1___conv2         getattr_G__net___layer1___1___conv2                         (getattr_g__net___layer1___1___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer1___1___bn2           getattr_G__net___layer1___1___bn2                           (getattr_g__net___layer1___1___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer1___1___relu_1        getattr_G__net___layer1___1___relu                          (getattr_g__net___layer1___1___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer1___1___conv3         getattr_G__net___layer1___1___conv3                         (getattr_g__net___layer1___1___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer1___1___bn3           getattr_G__net___layer1___1___bn3                           (getattr_g__net___layer1___1___conv3,)                                           {}\n",
      "call_function  iadd_1                                      <built-in function iadd>                                    (getattr_g__net___layer1___1___bn3, getattr_g__net___layer1___0___relu_2)        {}\n",
      "call_module    getattr_g__net___layer1___1___relu_2        getattr_G__net___layer1___1___relu                          (iadd_1,)                                                                        {}\n",
      "call_module    getattr_g__net___layer1___2___conv1         getattr_G__net___layer1___2___conv1                         (getattr_g__net___layer1___1___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer1___2___bn1           getattr_G__net___layer1___2___bn1                           (getattr_g__net___layer1___2___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer1___2___relu          getattr_G__net___layer1___2___relu                          (getattr_g__net___layer1___2___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer1___2___conv2         getattr_G__net___layer1___2___conv2                         (getattr_g__net___layer1___2___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer1___2___bn2           getattr_G__net___layer1___2___bn2                           (getattr_g__net___layer1___2___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer1___2___relu_1        getattr_G__net___layer1___2___relu                          (getattr_g__net___layer1___2___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer1___2___conv3         getattr_G__net___layer1___2___conv3                         (getattr_g__net___layer1___2___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer1___2___bn3           getattr_G__net___layer1___2___bn3                           (getattr_g__net___layer1___2___conv3,)                                           {}\n",
      "call_function  iadd_2                                      <built-in function iadd>                                    (getattr_g__net___layer1___2___bn3, getattr_g__net___layer1___1___relu_2)        {}\n",
      "call_module    getattr_g__net___layer1___2___relu_2        getattr_G__net___layer1___2___relu                          (iadd_2,)                                                                        {}\n",
      "call_module    getattr_g__net___layer2___0___conv1         getattr_G__net___layer2___0___conv1                         (getattr_g__net___layer1___2___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer2___0___bn1           getattr_G__net___layer2___0___bn1                           (getattr_g__net___layer2___0___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer2___0___relu          getattr_G__net___layer2___0___relu                          (getattr_g__net___layer2___0___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer2___0___conv2         getattr_G__net___layer2___0___conv2                         (getattr_g__net___layer2___0___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer2___0___bn2           getattr_G__net___layer2___0___bn2                           (getattr_g__net___layer2___0___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer2___0___relu_1        getattr_G__net___layer2___0___relu                          (getattr_g__net___layer2___0___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer2___0___conv3         getattr_G__net___layer2___0___conv3                         (getattr_g__net___layer2___0___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer2___0___bn3           getattr_G__net___layer2___0___bn3                           (getattr_g__net___layer2___0___conv3,)                                           {}\n",
      "call_module    getattr_g__net___layer2___0___downsample_0  getattr_G__net___layer2___0___downsample_0                  (getattr_g__net___layer1___2___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer2___0___downsample_1  getattr_G__net___layer2___0___downsample_1                  (getattr_g__net___layer2___0___downsample_0,)                                    {}\n",
      "call_function  iadd_3                                      <built-in function iadd>                                    (getattr_g__net___layer2___0___bn3, getattr_g__net___layer2___0___downsample_1)  {}\n",
      "call_module    getattr_g__net___layer2___0___relu_2        getattr_G__net___layer2___0___relu                          (iadd_3,)                                                                        {}\n",
      "call_module    getattr_g__net___layer2___1___conv1         getattr_G__net___layer2___1___conv1                         (getattr_g__net___layer2___0___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer2___1___bn1           getattr_G__net___layer2___1___bn1                           (getattr_g__net___layer2___1___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer2___1___relu          getattr_G__net___layer2___1___relu                          (getattr_g__net___layer2___1___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer2___1___conv2         getattr_G__net___layer2___1___conv2                         (getattr_g__net___layer2___1___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer2___1___bn2           getattr_G__net___layer2___1___bn2                           (getattr_g__net___layer2___1___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer2___1___relu_1        getattr_G__net___layer2___1___relu                          (getattr_g__net___layer2___1___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer2___1___conv3         getattr_G__net___layer2___1___conv3                         (getattr_g__net___layer2___1___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer2___1___bn3           getattr_G__net___layer2___1___bn3                           (getattr_g__net___layer2___1___conv3,)                                           {}\n",
      "call_function  iadd_4                                      <built-in function iadd>                                    (getattr_g__net___layer2___1___bn3, getattr_g__net___layer2___0___relu_2)        {}\n",
      "call_module    getattr_g__net___layer2___1___relu_2        getattr_G__net___layer2___1___relu                          (iadd_4,)                                                                        {}\n",
      "call_module    getattr_g__net___layer2___2___conv1         getattr_G__net___layer2___2___conv1                         (getattr_g__net___layer2___1___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer2___2___bn1           getattr_G__net___layer2___2___bn1                           (getattr_g__net___layer2___2___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer2___2___relu          getattr_G__net___layer2___2___relu                          (getattr_g__net___layer2___2___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer2___2___conv2         getattr_G__net___layer2___2___conv2                         (getattr_g__net___layer2___2___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer2___2___bn2           getattr_G__net___layer2___2___bn2                           (getattr_g__net___layer2___2___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer2___2___relu_1        getattr_G__net___layer2___2___relu                          (getattr_g__net___layer2___2___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer2___2___conv3         getattr_G__net___layer2___2___conv3                         (getattr_g__net___layer2___2___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer2___2___bn3           getattr_G__net___layer2___2___bn3                           (getattr_g__net___layer2___2___conv3,)                                           {}\n",
      "call_function  iadd_5                                      <built-in function iadd>                                    (getattr_g__net___layer2___2___bn3, getattr_g__net___layer2___1___relu_2)        {}\n",
      "call_module    getattr_g__net___layer2___2___relu_2        getattr_G__net___layer2___2___relu                          (iadd_5,)                                                                        {}\n",
      "call_module    getattr_g__net___layer2___3___conv1         getattr_G__net___layer2___3___conv1                         (getattr_g__net___layer2___2___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer2___3___bn1           getattr_G__net___layer2___3___bn1                           (getattr_g__net___layer2___3___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer2___3___relu          getattr_G__net___layer2___3___relu                          (getattr_g__net___layer2___3___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer2___3___conv2         getattr_G__net___layer2___3___conv2                         (getattr_g__net___layer2___3___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer2___3___bn2           getattr_G__net___layer2___3___bn2                           (getattr_g__net___layer2___3___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer2___3___relu_1        getattr_G__net___layer2___3___relu                          (getattr_g__net___layer2___3___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer2___3___conv3         getattr_G__net___layer2___3___conv3                         (getattr_g__net___layer2___3___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer2___3___bn3           getattr_G__net___layer2___3___bn3                           (getattr_g__net___layer2___3___conv3,)                                           {}\n",
      "call_function  iadd_6                                      <built-in function iadd>                                    (getattr_g__net___layer2___3___bn3, getattr_g__net___layer2___2___relu_2)        {}\n",
      "call_module    getattr_g__net___layer2___3___relu_2        getattr_G__net___layer2___3___relu                          (iadd_6,)                                                                        {}\n",
      "call_module    getattr_g__net___layer3___0___conv1         getattr_G__net___layer3___0___conv1                         (getattr_g__net___layer2___3___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer3___0___bn1           getattr_G__net___layer3___0___bn1                           (getattr_g__net___layer3___0___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer3___0___relu          getattr_G__net___layer3___0___relu                          (getattr_g__net___layer3___0___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer3___0___conv2         getattr_G__net___layer3___0___conv2                         (getattr_g__net___layer3___0___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer3___0___bn2           getattr_G__net___layer3___0___bn2                           (getattr_g__net___layer3___0___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer3___0___relu_1        getattr_G__net___layer3___0___relu                          (getattr_g__net___layer3___0___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer3___0___conv3         getattr_G__net___layer3___0___conv3                         (getattr_g__net___layer3___0___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer3___0___bn3           getattr_G__net___layer3___0___bn3                           (getattr_g__net___layer3___0___conv3,)                                           {}\n",
      "call_module    getattr_g__net___layer3___0___downsample_0  getattr_G__net___layer3___0___downsample_0                  (getattr_g__net___layer2___3___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer3___0___downsample_1  getattr_G__net___layer3___0___downsample_1                  (getattr_g__net___layer3___0___downsample_0,)                                    {}\n",
      "call_function  iadd_7                                      <built-in function iadd>                                    (getattr_g__net___layer3___0___bn3, getattr_g__net___layer3___0___downsample_1)  {}\n",
      "call_module    getattr_g__net___layer3___0___relu_2        getattr_G__net___layer3___0___relu                          (iadd_7,)                                                                        {}\n",
      "call_module    getattr_g__net___layer3___1___conv1         getattr_G__net___layer3___1___conv1                         (getattr_g__net___layer3___0___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer3___1___bn1           getattr_G__net___layer3___1___bn1                           (getattr_g__net___layer3___1___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer3___1___relu          getattr_G__net___layer3___1___relu                          (getattr_g__net___layer3___1___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer3___1___conv2         getattr_G__net___layer3___1___conv2                         (getattr_g__net___layer3___1___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer3___1___bn2           getattr_G__net___layer3___1___bn2                           (getattr_g__net___layer3___1___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer3___1___relu_1        getattr_G__net___layer3___1___relu                          (getattr_g__net___layer3___1___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer3___1___conv3         getattr_G__net___layer3___1___conv3                         (getattr_g__net___layer3___1___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer3___1___bn3           getattr_G__net___layer3___1___bn3                           (getattr_g__net___layer3___1___conv3,)                                           {}\n",
      "call_function  iadd_8                                      <built-in function iadd>                                    (getattr_g__net___layer3___1___bn3, getattr_g__net___layer3___0___relu_2)        {}\n",
      "call_module    getattr_g__net___layer3___1___relu_2        getattr_G__net___layer3___1___relu                          (iadd_8,)                                                                        {}\n",
      "call_module    getattr_g__net___layer3___2___conv1         getattr_G__net___layer3___2___conv1                         (getattr_g__net___layer3___1___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer3___2___bn1           getattr_G__net___layer3___2___bn1                           (getattr_g__net___layer3___2___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer3___2___relu          getattr_G__net___layer3___2___relu                          (getattr_g__net___layer3___2___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer3___2___conv2         getattr_G__net___layer3___2___conv2                         (getattr_g__net___layer3___2___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer3___2___bn2           getattr_G__net___layer3___2___bn2                           (getattr_g__net___layer3___2___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer3___2___relu_1        getattr_G__net___layer3___2___relu                          (getattr_g__net___layer3___2___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer3___2___conv3         getattr_G__net___layer3___2___conv3                         (getattr_g__net___layer3___2___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer3___2___bn3           getattr_G__net___layer3___2___bn3                           (getattr_g__net___layer3___2___conv3,)                                           {}\n",
      "call_function  iadd_9                                      <built-in function iadd>                                    (getattr_g__net___layer3___2___bn3, getattr_g__net___layer3___1___relu_2)        {}\n",
      "call_module    getattr_g__net___layer3___2___relu_2        getattr_G__net___layer3___2___relu                          (iadd_9,)                                                                        {}\n",
      "call_module    getattr_g__net___layer3___3___conv1         getattr_G__net___layer3___3___conv1                         (getattr_g__net___layer3___2___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer3___3___bn1           getattr_G__net___layer3___3___bn1                           (getattr_g__net___layer3___3___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer3___3___relu          getattr_G__net___layer3___3___relu                          (getattr_g__net___layer3___3___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer3___3___conv2         getattr_G__net___layer3___3___conv2                         (getattr_g__net___layer3___3___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer3___3___bn2           getattr_G__net___layer3___3___bn2                           (getattr_g__net___layer3___3___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer3___3___relu_1        getattr_G__net___layer3___3___relu                          (getattr_g__net___layer3___3___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer3___3___conv3         getattr_G__net___layer3___3___conv3                         (getattr_g__net___layer3___3___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer3___3___bn3           getattr_G__net___layer3___3___bn3                           (getattr_g__net___layer3___3___conv3,)                                           {}\n",
      "call_function  iadd_10                                     <built-in function iadd>                                    (getattr_g__net___layer3___3___bn3, getattr_g__net___layer3___2___relu_2)        {}\n",
      "call_module    getattr_g__net___layer3___3___relu_2        getattr_G__net___layer3___3___relu                          (iadd_10,)                                                                       {}\n",
      "call_module    getattr_g__net___layer3___4___conv1         getattr_G__net___layer3___4___conv1                         (getattr_g__net___layer3___3___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer3___4___bn1           getattr_G__net___layer3___4___bn1                           (getattr_g__net___layer3___4___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer3___4___relu          getattr_G__net___layer3___4___relu                          (getattr_g__net___layer3___4___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer3___4___conv2         getattr_G__net___layer3___4___conv2                         (getattr_g__net___layer3___4___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer3___4___bn2           getattr_G__net___layer3___4___bn2                           (getattr_g__net___layer3___4___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer3___4___relu_1        getattr_G__net___layer3___4___relu                          (getattr_g__net___layer3___4___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer3___4___conv3         getattr_G__net___layer3___4___conv3                         (getattr_g__net___layer3___4___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer3___4___bn3           getattr_G__net___layer3___4___bn3                           (getattr_g__net___layer3___4___conv3,)                                           {}\n",
      "call_function  iadd_11                                     <built-in function iadd>                                    (getattr_g__net___layer3___4___bn3, getattr_g__net___layer3___3___relu_2)        {}\n",
      "call_module    getattr_g__net___layer3___4___relu_2        getattr_G__net___layer3___4___relu                          (iadd_11,)                                                                       {}\n",
      "call_module    getattr_g__net___layer3___5___conv1         getattr_G__net___layer3___5___conv1                         (getattr_g__net___layer3___4___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer3___5___bn1           getattr_G__net___layer3___5___bn1                           (getattr_g__net___layer3___5___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer3___5___relu          getattr_G__net___layer3___5___relu                          (getattr_g__net___layer3___5___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer3___5___conv2         getattr_G__net___layer3___5___conv2                         (getattr_g__net___layer3___5___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer3___5___bn2           getattr_G__net___layer3___5___bn2                           (getattr_g__net___layer3___5___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer3___5___relu_1        getattr_G__net___layer3___5___relu                          (getattr_g__net___layer3___5___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer3___5___conv3         getattr_G__net___layer3___5___conv3                         (getattr_g__net___layer3___5___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer3___5___bn3           getattr_G__net___layer3___5___bn3                           (getattr_g__net___layer3___5___conv3,)                                           {}\n",
      "call_function  iadd_12                                     <built-in function iadd>                                    (getattr_g__net___layer3___5___bn3, getattr_g__net___layer3___4___relu_2)        {}\n",
      "call_module    getattr_g__net___layer3___5___relu_2        getattr_G__net___layer3___5___relu                          (iadd_12,)                                                                       {}\n",
      "call_module    getattr_g__net___layer4___0___conv1         getattr_G__net___layer4___0___conv1                         (getattr_g__net___layer3___5___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer4___0___bn1           getattr_G__net___layer4___0___bn1                           (getattr_g__net___layer4___0___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer4___0___relu          getattr_G__net___layer4___0___relu                          (getattr_g__net___layer4___0___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer4___0___conv2         getattr_G__net___layer4___0___conv2                         (getattr_g__net___layer4___0___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer4___0___bn2           getattr_G__net___layer4___0___bn2                           (getattr_g__net___layer4___0___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer4___0___relu_1        getattr_G__net___layer4___0___relu                          (getattr_g__net___layer4___0___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer4___0___conv3         getattr_G__net___layer4___0___conv3                         (getattr_g__net___layer4___0___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer4___0___bn3           getattr_G__net___layer4___0___bn3                           (getattr_g__net___layer4___0___conv3,)                                           {}\n",
      "call_module    getattr_g__net___layer4___0___downsample_0  getattr_G__net___layer4___0___downsample_0                  (getattr_g__net___layer3___5___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer4___0___downsample_1  getattr_G__net___layer4___0___downsample_1                  (getattr_g__net___layer4___0___downsample_0,)                                    {}\n",
      "call_function  iadd_13                                     <built-in function iadd>                                    (getattr_g__net___layer4___0___bn3, getattr_g__net___layer4___0___downsample_1)  {}\n",
      "call_module    getattr_g__net___layer4___0___relu_2        getattr_G__net___layer4___0___relu                          (iadd_13,)                                                                       {}\n",
      "call_module    getattr_g__net___layer4___1___conv1         getattr_G__net___layer4___1___conv1                         (getattr_g__net___layer4___0___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer4___1___bn1           getattr_G__net___layer4___1___bn1                           (getattr_g__net___layer4___1___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer4___1___relu          getattr_G__net___layer4___1___relu                          (getattr_g__net___layer4___1___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer4___1___conv2         getattr_G__net___layer4___1___conv2                         (getattr_g__net___layer4___1___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer4___1___bn2           getattr_G__net___layer4___1___bn2                           (getattr_g__net___layer4___1___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer4___1___relu_1        getattr_G__net___layer4___1___relu                          (getattr_g__net___layer4___1___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer4___1___conv3         getattr_G__net___layer4___1___conv3                         (getattr_g__net___layer4___1___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer4___1___bn3           getattr_G__net___layer4___1___bn3                           (getattr_g__net___layer4___1___conv3,)                                           {}\n",
      "call_function  iadd_14                                     <built-in function iadd>                                    (getattr_g__net___layer4___1___bn3, getattr_g__net___layer4___0___relu_2)        {}\n",
      "call_module    getattr_g__net___layer4___1___relu_2        getattr_G__net___layer4___1___relu                          (iadd_14,)                                                                       {}\n",
      "call_module    getattr_g__net___layer4___2___conv1         getattr_G__net___layer4___2___conv1                         (getattr_g__net___layer4___1___relu_2,)                                          {}\n",
      "call_module    getattr_g__net___layer4___2___bn1           getattr_G__net___layer4___2___bn1                           (getattr_g__net___layer4___2___conv1,)                                           {}\n",
      "call_module    getattr_g__net___layer4___2___relu          getattr_G__net___layer4___2___relu                          (getattr_g__net___layer4___2___bn1,)                                             {}\n",
      "call_module    getattr_g__net___layer4___2___conv2         getattr_G__net___layer4___2___conv2                         (getattr_g__net___layer4___2___relu,)                                            {}\n",
      "call_module    getattr_g__net___layer4___2___bn2           getattr_G__net___layer4___2___bn2                           (getattr_g__net___layer4___2___conv2,)                                           {}\n",
      "call_module    getattr_g__net___layer4___2___relu_1        getattr_G__net___layer4___2___relu                          (getattr_g__net___layer4___2___bn2,)                                             {}\n",
      "call_module    getattr_g__net___layer4___2___conv3         getattr_G__net___layer4___2___conv3                         (getattr_g__net___layer4___2___relu_1,)                                          {}\n",
      "call_module    getattr_g__net___layer4___2___bn3           getattr_G__net___layer4___2___bn3                           (getattr_g__net___layer4___2___conv3,)                                           {}\n",
      "call_function  iadd_15                                     <built-in function iadd>                                    (getattr_g__net___layer4___2___bn3, getattr_g__net___layer4___1___relu_2)        {}\n",
      "call_module    getattr_g__net___layer4___2___relu_2        getattr_G__net___layer4___2___relu                          (iadd_15,)                                                                       {}\n",
      "call_module    g__net___avgpool                            G__net___avgpool                                            (getattr_g__net___layer4___2___relu_2,)                                          {}\n",
      "call_function  flatten                                     <built-in method flatten of type object at 0x7ff353e9d100>  (g__net___avgpool, 1)                                                            {}\n",
      "call_module    g__net___fc                                 G__net___fc                                                 (flatten,)                                                                       {}\n",
      "call_module    criterion                                   criterion                                                   (g__net___fc, l_labels_)                                                         {}\n",
      "output         output                                      output                                                      ((criterion,),)                                                                  {}\n",
      "6.499421119689941\n"
     ]
    }
   ],
   "source": [
    "opt_model_segment = torch.compile(segment, backend=custom_backend)\n",
    "\n",
    "print(opt_model_segment(split_data_list[0], split_labels_list[0], losses=[]))"
   ]
  }
 ],
 "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.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}