From 853e088ed8c1f20f079f266ec552e85f9d6f1a9f Mon Sep 17 00:00:00 2001
From: Arnav Singhvi <arnavs11@berkeley.edu>
Date: Wed, 22 Feb 2023 12:39:46 -0800
Subject: [PATCH 1/5] network transformation function & testing for PyTorch to
 POET

---
 .../network_transform_pytorch.py              | 73 +++++++++++++++++++
 .../pytorch/test_resnet_network_transform.py  | 20 +++++
 .../pytorch/test_vgg_network_transform.py     | 16 ++++
 3 files changed, 109 insertions(+)
 create mode 100644 poet/architectures/network_transform_pytorch.py
 create mode 100644 poet/transformation_testing/pytorch/test_resnet_network_transform.py
 create mode 100644 poet/transformation_testing/pytorch/test_vgg_network_transform.py

diff --git a/poet/architectures/network_transform_pytorch.py b/poet/architectures/network_transform_pytorch.py
new file mode 100644
index 0000000..890fcb3
--- /dev/null
+++ b/poet/architectures/network_transform_pytorch.py
@@ -0,0 +1,73 @@
+from poet.power_computation import (
+    LinearLayer,
+    ReLULayer,
+    Conv2dLayer,
+    FlattenLayer,
+    TanHLayer,
+    SigmoidLayer,
+    SkipAddLayer,
+    DropoutLayer,
+    GradientLayer,
+    InputLayer,
+    SkipAddLayer,
+    CrossEntropyLoss,
+    GradientLayer,
+    BatchNorm2d,
+    MaxPool2d,
+    AvgPool2d,
+    GlobalAvgPool,
+    get_net_costs,
+)
+from poet.power_computation_transformer import QueryKeyValueMatrix, QKTMatrix, QKTVMatrix
+import torch.nn as nn
+import torchvision.models
+from torchvision.models.resnet import BasicBlock, Bottleneck
+
+
+# transforms input model's network layers to output graph with POET layers for PyTorch models
+def network_transform(net, layers, batch_size, num_classes, input_shape):
+    if isinstance(net, torchvision.models.resnet.ResNet):
+        modules = nn.Sequential(*list(net.children()))
+    elif isinstance(net, torchvision.models.vgg.VGG):
+        modules = list(net.children())
+    else:
+        modules = net
+    for module in modules:
+        if isinstance(module, nn.Sequential):
+            sequential_modules = [child for child in module]
+            input = network_transform(sequential_modules, [layers[-1]], batch_size, num_classes, input_shape)
+            layers.extend(input[1:])
+        if isinstance(module, BasicBlock) or isinstance(module, Bottleneck):
+            input = network_transform(nn.Sequential(*list(module.children())), [layers[-1]], batch_size, num_classes, input_shape)
+            layers.extend(input[1:])
+        if isinstance(module, nn.Linear):
+            lin_layer = LinearLayer(module.in_features, module.out_features, layers[-1])
+            act_layer = ReLULayer(lin_layer)
+            layers.extend([lin_layer, act_layer])
+        if isinstance(module, nn.ReLU):
+            relu_layer = ReLULayer(layers[-1])
+            layers.append(relu_layer)
+        if isinstance(module, nn.Conv2d):
+            conv_layer = Conv2dLayer(
+                module.in_channels, module.out_channels, module.kernel_size, module.stride[0], module.padding, layers[-1]
+            )
+            layers.append(conv_layer)
+        if isinstance(module, nn.BatchNorm2d):
+            layers.append(BatchNorm2d(layers[-1]))
+        if isinstance(module, nn.MaxPool2d):
+            layers.append(MaxPool2d((module.kernel_size, module.kernel_size), module.stride, layers[-1]))
+        if isinstance(module, nn.AvgPool2d):
+            layers.append(AvgPool2d(module.kernel_size, module.stride, layers[-1]))
+        if isinstance(module, nn.Tanh):
+            tanh_layer = TanHLayer(layers[-1])
+            layers.append(tanh_layer)
+        if isinstance(module, nn.Sigmoid):
+            sigmoid_layer = SigmoidLayer(layers[-1])
+            layers.append(sigmoid_layer)
+        if isinstance(module, nn.Flatten):
+            flatten_layer = FlattenLayer(layers[-1])
+            layers.append(flatten_layer)
+        if isinstance(module, nn.Dropout):
+            dropout_layer = DropoutLayer(layers[-1])
+            layers.append(dropout_layer)
+    return layers
diff --git a/poet/transformation_testing/pytorch/test_resnet_network_transform.py b/poet/transformation_testing/pytorch/test_resnet_network_transform.py
new file mode 100644
index 0000000..78537f7
--- /dev/null
+++ b/poet/transformation_testing/pytorch/test_resnet_network_transform.py
@@ -0,0 +1,20 @@
+from poet.power_computation import InputLayer
+import torchvision.models
+from poet.architectures.network_transform_pytorch import network_transform
+
+# transforms ResNet Model network layers into POET computation layers
+
+batch_size = (1,)
+input_shape = (3, 32, 32)
+num_classes = 10
+layers = [InputLayer((batch_size, *input_shape))]
+
+# comment out to output transformations for different ResNet models
+
+# #Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+# final_layers = network_transform(torchvision.models.resnet18(pretrained=True), layers, batch_size, num_classes, input_shape)
+# print(final_layers)
+
+# #Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+# final_layers = network_transform(torchvision.models.resnet50(pretrained=True), layers, batch_size, num_classes, input_shape)
+# print(final_layers)
diff --git a/poet/transformation_testing/pytorch/test_vgg_network_transform.py b/poet/transformation_testing/pytorch/test_vgg_network_transform.py
new file mode 100644
index 0000000..9f753cf
--- /dev/null
+++ b/poet/transformation_testing/pytorch/test_vgg_network_transform.py
@@ -0,0 +1,16 @@
+from poet.power_computation import InputLayer
+import torch
+from poet.architectures.network_transform_pytorch import network_transform
+
+# transforms VGG Model network layers into POET computation layers
+
+batch_size = (1,)
+input_shape = (3, 32, 32)
+num_classes = 10
+layers = [InputLayer((batch_size, *input_shape))]
+
+# VGG16 model transformation - https://download.pytorch.org/models/vgg16-397923af.pth
+final_layers = network_transform(
+    torch.hub.load("pytorch/vision:v0.10.0", "vgg16", pretrained=True), layers, batch_size, num_classes, input_shape
+)
+print(final_layers)

From 06349a8c876540cb3aa5a31f2094214eb2595ac5 Mon Sep 17 00:00:00 2001
From: Arnav Singhvi <arnavs11@berkeley.edu>
Date: Wed, 22 Feb 2023 12:46:58 -0800
Subject: [PATCH 2/5] network transformation function & testing for TensorFlow
 to POET

---
 .../network_transform_tensorflow.py           | 62 ++++++++++++++
 .../test_resnet_network_transform.py          | 81 +++++++++++++++++++
 2 files changed, 143 insertions(+)
 create mode 100644 poet/architectures/network_transform_tensorflow.py
 create mode 100644 poet/transformation_testing/tensorflow/test_resnet_network_transform.py

diff --git a/poet/architectures/network_transform_tensorflow.py b/poet/architectures/network_transform_tensorflow.py
new file mode 100644
index 0000000..4ef9596
--- /dev/null
+++ b/poet/architectures/network_transform_tensorflow.py
@@ -0,0 +1,62 @@
+from poet.power_computation import (
+    LinearLayer,
+    ReLULayer,
+    Conv2dLayer,
+    FlattenLayer,
+    TanHLayer,
+    SigmoidLayer,
+    SkipAddLayer,
+    DropoutLayer,
+    GradientLayer,
+    InputLayer,
+    SkipAddLayer,
+    CrossEntropyLoss,
+    GradientLayer,
+    BatchNorm2d,
+    MaxPool2d,
+    AvgPool2d,
+    GlobalAvgPool,
+    get_net_costs,
+)
+from poet.power_computation_transformer import QueryKeyValueMatrix, QKTMatrix, QKTVMatrix
+from torchvision.models.resnet import BasicBlock, Bottleneck
+import tensorflow as tf
+
+
+##transforms input model's network layers to output graph with POET layers for TensorFlow models
+def network_transform(net, layers, batch_size, num_classes, input_shape):
+    for module in net:
+        if isinstance(module, tf.keras.layers.Dense):
+            lin_layer = LinearLayer(module.units, module.units, layers[-1])
+            act_layer = ReLULayer(lin_layer)
+            layers.extend([lin_layer, act_layer])
+        if isinstance(module, tf.keras.layers.Activation) and module._name == "relu":
+            relu_layer = ReLULayer(layers[-1])
+            layers.append(relu_layer)
+        if isinstance(module, tf.keras.layers.Conv2D):
+            if module.padding == "valid":
+                padding = (0, 0)
+            elif module.padding == "same":
+                padding = (1, 1)
+            conv_layer = Conv2dLayer(1, module.filters, module.kernel_size, module.strides[0], padding, layers[-1])
+            layers.append(conv_layer)
+        if isinstance(module, tf.keras.layers.BatchNormalization):
+            layers.append(BatchNorm2d(layers[-1]))
+        if isinstance(module, tf.keras.layers.MaxPool2D):
+            layers.append(MaxPool2d(module.pool_size, module.strides[0], layers[-1]))
+        if isinstance(module, tf.keras.layers.GlobalAveragePooling2D):
+            if module.keepdims:
+                layers.append(GlobalAvgPool(layers[-1]))
+        if isinstance(module, tf.keras.layers.Activation) and module._name == "tanh":
+            tanh_layer = TanHLayer(layers[-1])
+            layers.append(tanh_layer)
+        if isinstance(module, tf.keras.layers.Activation) and module._name == "sigmoid":
+            sigmoid_layer = SigmoidLayer(layers[-1])
+            layers.append(sigmoid_layer)
+        if isinstance(module, tf.keras.layers.Flatten):
+            flatten_layer = FlattenLayer(layers[-1])
+            layers.append(flatten_layer)
+        if isinstance(module, tf.keras.layers.Dropout):
+            dropout_layer = DropoutLayer(layers[-1])
+            layers.append(dropout_layer)
+    return layers
diff --git a/poet/transformation_testing/tensorflow/test_resnet_network_transform.py b/poet/transformation_testing/tensorflow/test_resnet_network_transform.py
new file mode 100644
index 0000000..501e75b
--- /dev/null
+++ b/poet/transformation_testing/tensorflow/test_resnet_network_transform.py
@@ -0,0 +1,81 @@
+from poet.architectures.network_transform_tensorflow import network_transform
+from poet.power_computation import InputLayer
+import tensorflow as tf
+
+# transforms VGG Model network layers into POET computation layers
+
+
+# defining ResNet model in TensorFlow
+def BasicBlock(inputs, num_channels, kernel_size, num_blocks, skip_blocks, name):
+    """Basic residual block"""
+    x = inputs
+    for i in range(num_blocks):
+        if i not in skip_blocks:
+            x1 = ConvNormRelu(x, num_channels, kernel_size, strides=[1, 1], name=name + "." + str(i))
+            x = tf.keras.layers.Add()([x, x1])
+            x = tf.keras.layers.Activation("relu")(x)
+    return x
+
+
+def BasicBlockDown(inputs, num_channels, kernel_size, name):
+    """Residual block with strided downsampling"""
+    x = inputs
+    x1 = ConvNormRelu(x, num_channels, kernel_size, strides=[2, 1], name=name + ".0")
+    x = tf.keras.layers.Conv2D(
+        num_channels, kernel_size=1, strides=2, padding="same", activation="linear", use_bias=False, name=name + ".0.downsample.0"
+    )(x)
+    x = tf.keras.layers.BatchNormalization(momentum=0.1, epsilon=1e-5, name=name + ".0.downsample.1")(x)
+    x = tf.keras.layers.Add()([x, x1])
+    x = tf.keras.layers.Activation("relu")(x)
+    return x
+
+
+def ConvNormRelu(x, num_channels, kernel_size, strides, name):
+    """Layer consisting of 2 consecutive batch normalizations with 1 first relu"""
+    if strides[0] == 2:
+        x = tf.keras.layers.ZeroPadding2D(padding=(1, 1), name=name + ".pad")(x)
+        x = tf.keras.layers.Conv2D(
+            num_channels, kernel_size, strides[0], padding="valid", activation="linear", use_bias=False, name=name + ".conv1"
+        )(x)
+    else:
+        x = tf.keras.layers.Conv2D(
+            num_channels, kernel_size, strides[0], padding="same", activation="linear", use_bias=False, name=name + ".conv1"
+        )(x)
+    x = tf.keras.layers.BatchNormalization(momentum=0.1, epsilon=1e-5, name=name + ".bn1")(x)
+    x = tf.keras.layers.Activation("relu")(x)
+    x = tf.keras.layers.Conv2D(
+        num_channels, kernel_size, strides[1], padding="same", activation="linear", use_bias=False, name=name + ".conv2"
+    )(x)
+    x = tf.keras.layers.BatchNormalization(momentum=0.1, epsilon=1e-5, name=name + ".bn2")(x)
+    return x
+
+
+def ResNet18(inputs):
+    x = tf.keras.layers.ZeroPadding2D(padding=(3, 3), name="pad")(inputs)
+    x = tf.keras.layers.Conv2D(filters=64, kernel_size=7, strides=2, padding="valid", activation="linear", use_bias=False, name="conv1")(x)
+    x = tf.keras.layers.BatchNormalization(momentum=0.1, epsilon=1e-5, name="bn1")(x)
+    x = tf.keras.layers.Activation("relu", name="relu")(x)
+    x = tf.keras.layers.ZeroPadding2D(padding=(1, 1), name="pad1")(x)
+    x = tf.keras.layers.MaxPool2D(pool_size=3, strides=2, padding="valid", name="maxpool")(x)
+    x = BasicBlock(x, num_channels=64, kernel_size=3, num_blocks=2, skip_blocks=[], name="layer1")
+    x = BasicBlockDown(x, num_channels=128, kernel_size=3, name="layer2")
+    x = BasicBlock(x, num_channels=128, kernel_size=3, num_blocks=2, skip_blocks=[0], name="layer2")
+    x = BasicBlockDown(x, num_channels=256, kernel_size=3, name="layer3")
+    x = BasicBlock(x, num_channels=256, kernel_size=3, num_blocks=2, skip_blocks=[0], name="layer3")
+    x = BasicBlockDown(x, num_channels=512, kernel_size=3, name="layer4")
+    x = BasicBlock(x, num_channels=512, kernel_size=3, num_blocks=2, skip_blocks=[0], name="layer4")
+    x = tf.keras.layers.GlobalAveragePooling2D(name="avgpool")(x)
+    x = tf.keras.layers.Dense(units=1000, use_bias=True, activation="linear", name="fc")(x)
+    return x
+
+
+inputs = tf.keras.Input((None, None, 3))
+resnet_tf = ResNet18(inputs)
+model = tf.keras.Model(inputs, resnet_tf)
+batch_size = (1,)
+input_shape = (3, 32, 32)
+num_classes = 10
+layers = [InputLayer((batch_size, *input_shape))]
+
+final_layers = network_transform([layer for layer in model.layers], layers, batch_size, num_classes, input_shape)
+print(final_layers)

From 305f86156ddc0eec21b14ad9bf073e59263e5d8b Mon Sep 17 00:00:00 2001
From: Arnav Singhvi <arnavs11@berkeley.edu>
Date: Wed, 22 Feb 2023 12:51:00 -0800
Subject: [PATCH 3/5] graph transformation function for converting PyTorch
 model graph to POET layers graph

---
 poet/architectures/graph_transformation.py    | 59 +++++++++++++++++++
 .../pytorch/test_resnet_graph_transform.py    | 19 ++++++
 2 files changed, 78 insertions(+)
 create mode 100644 poet/architectures/graph_transformation.py
 create mode 100644 poet/transformation_testing/pytorch/test_resnet_graph_transform.py

diff --git a/poet/architectures/graph_transformation.py b/poet/architectures/graph_transformation.py
new file mode 100644
index 0000000..f6daad2
--- /dev/null
+++ b/poet/architectures/graph_transformation.py
@@ -0,0 +1,59 @@
+from poet.power_computation import (
+    LinearLayer,
+    ReLULayer,
+    Conv2dLayer,
+    FlattenLayer,
+    TanHLayer,
+    SigmoidLayer,
+    SkipAddLayer,
+    DropoutLayer,
+    GradientLayer,
+    InputLayer,
+    SkipAddLayer,
+    CrossEntropyLoss,
+    GradientLayer,
+    BatchNorm2d,
+    MaxPool2d,
+    AvgPool2d,
+    GlobalAvgPool,
+    get_net_costs,
+)
+
+
+# transforms input model's graph to output graph with POET layer nodes
+def graph_transform(traced):
+    for n in traced.graph.nodes:
+        if "<built-in function" in str(n.target):
+            continue
+        elif "fc" in str(n.target):
+            with traced.graph.inserting_after(n):
+                new_node = traced.graph.call_function(LinearLayer, n.args, n.kwargs)
+                n.replace_all_uses_with(new_node)
+            traced.graph.erase_node(n)
+        elif "flatten" in str(n.target):
+            with traced.graph.inserting_after(n):
+                new_node = traced.graph.call_function(FlattenLayer, n.args, n.kwargs)
+                n.replace_all_uses_with(new_node)
+            traced.graph.erase_node(n)
+        elif "relu" in str(n.target):
+            with traced.graph.inserting_after(n):
+                new_node = traced.graph.call_function(ReLULayer, n.args, n.kwargs)
+                n.replace_all_uses_with(new_node)
+            traced.graph.erase_node(n)
+        elif "conv" in str(n.target):
+            with traced.graph.inserting_after(n):
+                new_node = traced.graph.call_function(Conv2dLayer, n.args, n.kwargs)
+                n.replace_all_uses_with(new_node)
+            traced.graph.erase_node(n)
+        elif "bn" in str(n.target):
+            with traced.graph.inserting_after(n):
+                new_node = traced.graph.call_function(BatchNorm2d, n.args, n.kwargs)
+                n.replace_all_uses_with(new_node)
+            traced.graph.erase_node(n)
+        elif "maxpool" in str(n.target):
+            with traced.graph.inserting_after(n):
+                new_node = traced.graph.call_function(MaxPool2d, n.args, n.kwargs)
+                n.replace_all_uses_with(new_node)
+            traced.graph.erase_node(n)
+    traced.recompile()
+    return traced
diff --git a/poet/transformation_testing/pytorch/test_resnet_graph_transform.py b/poet/transformation_testing/pytorch/test_resnet_graph_transform.py
new file mode 100644
index 0000000..194b4d5
--- /dev/null
+++ b/poet/transformation_testing/pytorch/test_resnet_graph_transform.py
@@ -0,0 +1,19 @@
+from torch.fx import symbolic_trace
+import torchvision
+from poet.architectures.graph_transformation import graph_transform
+
+# transforms ResNet Model graph into POET layers nodes
+
+# #Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+# traced = symbolic_trace(torchvision.models.resnet18(pretrained=True))
+# poet_traced = graph_transform(traced)
+# for n in poet_traced.graph.nodes:
+#   print(n.target)
+#   print(n.name)
+
+# #Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+# traced = symbolic_trace(torchvision.models.resnet50(pretrained=True))
+# poet_traced = graph_transform(traced)
+# for n in poet_traced.graph.nodes:
+#   print(n.target)
+#   print(n.name)

From ce1c44da8115498fbaa15c5d58e54bb79fc66b99 Mon Sep 17 00:00:00 2001
From: Arnav Singhvi <arnavs11@berkeley.edu>
Date: Thu, 2 Mar 2023 16:34:33 -0800
Subject: [PATCH 4/5] updated typing and graph transformation edge case,
 revised changes from pull request comments

---
 poet/architectures/graph_transformation.py    | 13 ++++++++--
 .../network_transform_pytorch.py              |  5 ++--
 .../pytorch/test_resnet_graph_transform.py    | 25 +++++++++----------
 .../pytorch/test_resnet_network_transform.py  | 14 +++++------
 .../test_resnet_network_transform.py          |  4 +--
 5 files changed, 33 insertions(+), 28 deletions(-)

diff --git a/poet/architectures/graph_transformation.py b/poet/architectures/graph_transformation.py
index f6daad2..06abbcd 100644
--- a/poet/architectures/graph_transformation.py
+++ b/poet/architectures/graph_transformation.py
@@ -1,3 +1,4 @@
+import torch
 from poet.power_computation import (
     LinearLayer,
     ReLULayer,
@@ -21,9 +22,11 @@
 
 
 # transforms input model's graph to output graph with POET layer nodes
-def graph_transform(traced):
+def graph_transform(traced: torch.fx.graph_module.GraphModule) -> torch.fx.graph_module.GraphModule: 
     for n in traced.graph.nodes:
-        if "<built-in function" in str(n.target):
+        #ignores built-in functions and input x which are not layer nodes in the model graph
+        #ignores poet layer nodes which are added to the model graph from this function
+        if "<built-in function" in str(n.target) or "poet" in str(n.target) or "x" == str(n.target):
             continue
         elif "fc" in str(n.target):
             with traced.graph.inserting_after(n):
@@ -55,5 +58,11 @@ def graph_transform(traced):
                 new_node = traced.graph.call_function(MaxPool2d, n.args, n.kwargs)
                 n.replace_all_uses_with(new_node)
             traced.graph.erase_node(n)
+        else:
+            user_input = input(str(n.target) + ' is not supported by POET layers. Would you like to proceed? (y/n)')
+            if user_input.lower() == 'y':
+                continue
+            else:
+                exit(0)
     traced.recompile()
     return traced
diff --git a/poet/architectures/network_transform_pytorch.py b/poet/architectures/network_transform_pytorch.py
index 890fcb3..105118d 100644
--- a/poet/architectures/network_transform_pytorch.py
+++ b/poet/architectures/network_transform_pytorch.py
@@ -18,7 +18,6 @@
     GlobalAvgPool,
     get_net_costs,
 )
-from poet.power_computation_transformer import QueryKeyValueMatrix, QKTMatrix, QKTVMatrix
 import torch.nn as nn
 import torchvision.models
 from torchvision.models.resnet import BasicBlock, Bottleneck
@@ -39,7 +38,7 @@ def network_transform(net, layers, batch_size, num_classes, input_shape):
             layers.extend(input[1:])
         if isinstance(module, BasicBlock) or isinstance(module, Bottleneck):
             input = network_transform(nn.Sequential(*list(module.children())), [layers[-1]], batch_size, num_classes, input_shape)
-            layers.extend(input[1:])
+            layers.extend(input[1:])        
         if isinstance(module, nn.Linear):
             lin_layer = LinearLayer(module.in_features, module.out_features, layers[-1])
             act_layer = ReLULayer(lin_layer)
@@ -70,4 +69,4 @@ def network_transform(net, layers, batch_size, num_classes, input_shape):
         if isinstance(module, nn.Dropout):
             dropout_layer = DropoutLayer(layers[-1])
             layers.append(dropout_layer)
-    return layers
+    return layers
\ No newline at end of file
diff --git a/poet/transformation_testing/pytorch/test_resnet_graph_transform.py b/poet/transformation_testing/pytorch/test_resnet_graph_transform.py
index 194b4d5..a441003 100644
--- a/poet/transformation_testing/pytorch/test_resnet_graph_transform.py
+++ b/poet/transformation_testing/pytorch/test_resnet_graph_transform.py
@@ -1,19 +1,18 @@
 from torch.fx import symbolic_trace
 import torchvision
 from poet.architectures.graph_transformation import graph_transform
-
 # transforms ResNet Model graph into POET layers nodes
 
-# #Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
-# traced = symbolic_trace(torchvision.models.resnet18(pretrained=True))
-# poet_traced = graph_transform(traced)
-# for n in poet_traced.graph.nodes:
-#   print(n.target)
-#   print(n.name)
+# Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+traced = symbolic_trace(torchvision.models.resnet18(pretrained=True))
+poet_traced = graph_transform(traced)
+for n in poet_traced.graph.nodes:
+  print(n.target)
+  print(n.name)
 
-# #Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
-# traced = symbolic_trace(torchvision.models.resnet50(pretrained=True))
-# poet_traced = graph_transform(traced)
-# for n in poet_traced.graph.nodes:
-#   print(n.target)
-#   print(n.name)
+#Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+traced = symbolic_trace(torchvision.models.resnet50(pretrained=True))
+poet_traced = graph_transform(traced)
+for n in poet_traced.graph.nodes:
+  print(n.target)
+  print(n.name)
diff --git a/poet/transformation_testing/pytorch/test_resnet_network_transform.py b/poet/transformation_testing/pytorch/test_resnet_network_transform.py
index 78537f7..87e5624 100644
--- a/poet/transformation_testing/pytorch/test_resnet_network_transform.py
+++ b/poet/transformation_testing/pytorch/test_resnet_network_transform.py
@@ -9,12 +9,10 @@
 num_classes = 10
 layers = [InputLayer((batch_size, *input_shape))]
 
-# comment out to output transformations for different ResNet models
+#Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+final_layers = network_transform(torchvision.models.resnet18(pretrained=True), layers, batch_size, num_classes, input_shape)
+print(final_layers)
 
-# #Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
-# final_layers = network_transform(torchvision.models.resnet18(pretrained=True), layers, batch_size, num_classes, input_shape)
-# print(final_layers)
-
-# #Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
-# final_layers = network_transform(torchvision.models.resnet50(pretrained=True), layers, batch_size, num_classes, input_shape)
-# print(final_layers)
+#Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+final_layers = network_transform(torchvision.models.resnet50(pretrained=True), layers, batch_size, num_classes, input_shape)
+print(final_layers)
diff --git a/poet/transformation_testing/tensorflow/test_resnet_network_transform.py b/poet/transformation_testing/tensorflow/test_resnet_network_transform.py
index 501e75b..3870c68 100644
--- a/poet/transformation_testing/tensorflow/test_resnet_network_transform.py
+++ b/poet/transformation_testing/tensorflow/test_resnet_network_transform.py
@@ -2,7 +2,7 @@
 from poet.power_computation import InputLayer
 import tensorflow as tf
 
-# transforms VGG Model network layers into POET computation layers
+# transforms ResNet Model network layers into POET computation layers
 
 
 # defining ResNet model in TensorFlow
@@ -78,4 +78,4 @@ def ResNet18(inputs):
 layers = [InputLayer((batch_size, *input_shape))]
 
 final_layers = network_transform([layer for layer in model.layers], layers, batch_size, num_classes, input_shape)
-print(final_layers)
+print(final_layers)
\ No newline at end of file

From 6dbcad4bb130059f173cce64e350a4dbd949a0cf Mon Sep 17 00:00:00 2001
From: Arnav Singhvi <arnav11.singhvi@gmail.com>
Date: Tue, 4 Apr 2023 10:15:02 -0700
Subject: [PATCH 5/5] renamed to framework_integration, adding testing files to
 test folder

---
 poet/architectures/graph_transformation.py    | 10 +--
 .../network_transform_pytorch.py              |  4 +-
 .../pytorch/test_resnet_graph_transform.py    | 11 +--
 .../pytorch/test_resnet_network_transform.py  |  4 +-
 .../pytorch/test_vgg_network_transform.py     |  0
 .../test_resnet_network_transform.py          |  2 +-
 poet/poet_solver.py                           |  2 +-
 poet/poet_solver_gurobi.py                    |  3 +-
 poet/utils/checkmate/core/utils/graph.py      |  2 +-
 test/test_pytorch_resnet_graph_transform.py   | 21 +++++
 test/test_pytorch_resnet_network_transform.py | 20 +++++
 test/test_pytorch_vgg_network_transform.py    | 17 ++++
 ...est_tensorflow_resnet_network_transform.py | 82 +++++++++++++++++++
 13 files changed, 160 insertions(+), 18 deletions(-)
 rename poet/{transformation_testing => framework_integration}/pytorch/test_resnet_graph_transform.py (72%)
 rename poet/{transformation_testing => framework_integration}/pytorch/test_resnet_network_transform.py (66%)
 rename poet/{transformation_testing => framework_integration}/pytorch/test_vgg_network_transform.py (100%)
 rename poet/{transformation_testing => framework_integration}/tensorflow/test_resnet_network_transform.py (99%)
 create mode 100644 test/test_pytorch_resnet_graph_transform.py
 create mode 100644 test/test_pytorch_resnet_network_transform.py
 create mode 100644 test/test_pytorch_vgg_network_transform.py
 create mode 100644 test/test_tensorflow_resnet_network_transform.py

diff --git a/poet/architectures/graph_transformation.py b/poet/architectures/graph_transformation.py
index 06abbcd..8d76398 100644
--- a/poet/architectures/graph_transformation.py
+++ b/poet/architectures/graph_transformation.py
@@ -22,10 +22,10 @@
 
 
 # transforms input model's graph to output graph with POET layer nodes
-def graph_transform(traced: torch.fx.graph_module.GraphModule) -> torch.fx.graph_module.GraphModule: 
+def graph_transform(traced: torch.fx.graph_module.GraphModule) -> torch.fx.graph_module.GraphModule:
     for n in traced.graph.nodes:
-        #ignores built-in functions and input x which are not layer nodes in the model graph
-        #ignores poet layer nodes which are added to the model graph from this function
+        # ignores built-in functions and input x which are not layer nodes in the model graph
+        # ignores poet layer nodes which are added to the model graph from this function
         if "<built-in function" in str(n.target) or "poet" in str(n.target) or "x" == str(n.target):
             continue
         elif "fc" in str(n.target):
@@ -59,8 +59,8 @@ def graph_transform(traced: torch.fx.graph_module.GraphModule) -> torch.fx.graph
                 n.replace_all_uses_with(new_node)
             traced.graph.erase_node(n)
         else:
-            user_input = input(str(n.target) + ' is not supported by POET layers. Would you like to proceed? (y/n)')
-            if user_input.lower() == 'y':
+            user_input = input(str(n.target) + " is not supported by POET layers. Would you like to proceed? (y/n)")
+            if user_input.lower() == "y":
                 continue
             else:
                 exit(0)
diff --git a/poet/architectures/network_transform_pytorch.py b/poet/architectures/network_transform_pytorch.py
index 105118d..b7be455 100644
--- a/poet/architectures/network_transform_pytorch.py
+++ b/poet/architectures/network_transform_pytorch.py
@@ -38,7 +38,7 @@ def network_transform(net, layers, batch_size, num_classes, input_shape):
             layers.extend(input[1:])
         if isinstance(module, BasicBlock) or isinstance(module, Bottleneck):
             input = network_transform(nn.Sequential(*list(module.children())), [layers[-1]], batch_size, num_classes, input_shape)
-            layers.extend(input[1:])        
+            layers.extend(input[1:])
         if isinstance(module, nn.Linear):
             lin_layer = LinearLayer(module.in_features, module.out_features, layers[-1])
             act_layer = ReLULayer(lin_layer)
@@ -69,4 +69,4 @@ def network_transform(net, layers, batch_size, num_classes, input_shape):
         if isinstance(module, nn.Dropout):
             dropout_layer = DropoutLayer(layers[-1])
             layers.append(dropout_layer)
-    return layers
\ No newline at end of file
+    return layers
diff --git a/poet/transformation_testing/pytorch/test_resnet_graph_transform.py b/poet/framework_integration/pytorch/test_resnet_graph_transform.py
similarity index 72%
rename from poet/transformation_testing/pytorch/test_resnet_graph_transform.py
rename to poet/framework_integration/pytorch/test_resnet_graph_transform.py
index a441003..4decf40 100644
--- a/poet/transformation_testing/pytorch/test_resnet_graph_transform.py
+++ b/poet/framework_integration/pytorch/test_resnet_graph_transform.py
@@ -1,18 +1,19 @@
 from torch.fx import symbolic_trace
 import torchvision
 from poet.architectures.graph_transformation import graph_transform
+
 # transforms ResNet Model graph into POET layers nodes
 
 # Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
 traced = symbolic_trace(torchvision.models.resnet18(pretrained=True))
 poet_traced = graph_transform(traced)
 for n in poet_traced.graph.nodes:
-  print(n.target)
-  print(n.name)
+    print(n.target)
+    print(n.name)
 
-#Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+# Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
 traced = symbolic_trace(torchvision.models.resnet50(pretrained=True))
 poet_traced = graph_transform(traced)
 for n in poet_traced.graph.nodes:
-  print(n.target)
-  print(n.name)
+    print(n.target)
+    print(n.name)
diff --git a/poet/transformation_testing/pytorch/test_resnet_network_transform.py b/poet/framework_integration/pytorch/test_resnet_network_transform.py
similarity index 66%
rename from poet/transformation_testing/pytorch/test_resnet_network_transform.py
rename to poet/framework_integration/pytorch/test_resnet_network_transform.py
index 87e5624..8cea2c0 100644
--- a/poet/transformation_testing/pytorch/test_resnet_network_transform.py
+++ b/poet/framework_integration/pytorch/test_resnet_network_transform.py
@@ -9,10 +9,10 @@
 num_classes = 10
 layers = [InputLayer((batch_size, *input_shape))]
 
-#Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+# Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
 final_layers = network_transform(torchvision.models.resnet18(pretrained=True), layers, batch_size, num_classes, input_shape)
 print(final_layers)
 
-#Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+# Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
 final_layers = network_transform(torchvision.models.resnet50(pretrained=True), layers, batch_size, num_classes, input_shape)
 print(final_layers)
diff --git a/poet/transformation_testing/pytorch/test_vgg_network_transform.py b/poet/framework_integration/pytorch/test_vgg_network_transform.py
similarity index 100%
rename from poet/transformation_testing/pytorch/test_vgg_network_transform.py
rename to poet/framework_integration/pytorch/test_vgg_network_transform.py
diff --git a/poet/transformation_testing/tensorflow/test_resnet_network_transform.py b/poet/framework_integration/tensorflow/test_resnet_network_transform.py
similarity index 99%
rename from poet/transformation_testing/tensorflow/test_resnet_network_transform.py
rename to poet/framework_integration/tensorflow/test_resnet_network_transform.py
index 3870c68..6a469d1 100644
--- a/poet/transformation_testing/tensorflow/test_resnet_network_transform.py
+++ b/poet/framework_integration/tensorflow/test_resnet_network_transform.py
@@ -78,4 +78,4 @@ def ResNet18(inputs):
 layers = [InputLayer((batch_size, *input_shape))]
 
 final_layers = network_transform([layer for layer in model.layers], layers, batch_size, num_classes, input_shape)
-print(final_layers)
\ No newline at end of file
+print(final_layers)
diff --git a/poet/poet_solver.py b/poet/poet_solver.py
index 8cff292..15976f4 100644
--- a/poet/poet_solver.py
+++ b/poet/poet_solver.py
@@ -107,7 +107,7 @@ def _initialize_variables(self):
 
     def _create_correctness_constraints(self):
         # ensure all computations are possible
-        for (u, v) in self.g.edge_list:
+        for u, v in self.g.edge_list:
             for t in range(self.T):
                 self.m += self.R[t][v] <= self.R[t][u] + self.SRam[t][u]
         # ensure all checkpoints are in memory
diff --git a/poet/poet_solver_gurobi.py b/poet/poet_solver_gurobi.py
index cb25b29..e98da9b 100644
--- a/poet/poet_solver_gurobi.py
+++ b/poet/poet_solver_gurobi.py
@@ -13,6 +13,7 @@
 
 # noinspection PyPackageRequirements
 
+
 # POET ILP defined using Gurobi
 class POETSolverGurobi:
     def __init__(
@@ -124,7 +125,7 @@ def _disable_paging(self):
 
     def _create_correctness_constraints(self):
         # ensure all computations are possible
-        for (u, v) in self.g.edge_list:
+        for u, v in self.g.edge_list:
             for t in range(self.T):
                 self.m.addLConstr(self.R[t, v], GRB.LESS_EQUAL, self.R[t, u] + self.SRam[t, u])
         # ensure all checkpoints are in memory
diff --git a/poet/utils/checkmate/core/utils/graph.py b/poet/utils/checkmate/core/utils/graph.py
index 539baea..0a3e084 100644
--- a/poet/utils/checkmate/core/utils/graph.py
+++ b/poet/utils/checkmate/core/utils/graph.py
@@ -10,7 +10,7 @@
 def edge_to_adj_list(E: EdgeList, convert_undirected=False):
     """Returns an (undirected / bidirectional) adjacency list"""
     adj_list = defaultdict(set)
-    for (i, j) in E:
+    for i, j in E:
         adj_list[i].add(j)
         if convert_undirected:
             adj_list[j].add(i)
diff --git a/test/test_pytorch_resnet_graph_transform.py b/test/test_pytorch_resnet_graph_transform.py
new file mode 100644
index 0000000..0b3d6bb
--- /dev/null
+++ b/test/test_pytorch_resnet_graph_transform.py
@@ -0,0 +1,21 @@
+from torch.fx import symbolic_trace
+import torchvision
+from poet.architectures.graph_transformation import graph_transform
+
+# transforms ResNet Model graph into POET layers nodes
+
+# Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+traced = symbolic_trace(torchvision.models.resnet18(pretrained=True))
+poet_traced = graph_transform(traced)
+for n in poet_traced.graph.nodes:
+    print(n.target)
+    print(n.name)
+
+
+# Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+traced = symbolic_trace(torchvision.models.resnet50(pretrained=True))
+poet_traced = graph_transform(traced)
+for n in poet_traced.graph.nodes:
+    print(n.target)
+    print(n.name)
+assert n.target == "output" and n.name == "output"
diff --git a/test/test_pytorch_resnet_network_transform.py b/test/test_pytorch_resnet_network_transform.py
new file mode 100644
index 0000000..f7ca79f
--- /dev/null
+++ b/test/test_pytorch_resnet_network_transform.py
@@ -0,0 +1,20 @@
+from poet.power_computation import InputLayer
+import torchvision.models
+from poet.architectures.network_transform_pytorch import network_transform
+
+# transforms ResNet Model network layers into POET computation layers
+
+batch_size = (1,)
+input_shape = (3, 32, 32)
+num_classes = 10
+layers = [InputLayer((batch_size, *input_shape))]
+
+# Resnet18 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+final_layers = network_transform(torchvision.models.resnet18(pretrained=True), layers, batch_size, num_classes, input_shape)
+print(final_layers)
+assert len(final_layers) > 0
+
+# Resnet50 model transformation - https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py, commit: 7dc5e5bd60b55eb4e6ea5c1265d6dc7b17d2e917
+final_layers = network_transform(torchvision.models.resnet50(pretrained=True), layers, batch_size, num_classes, input_shape)
+print(final_layers)
+assert len(final_layers) > 0
diff --git a/test/test_pytorch_vgg_network_transform.py b/test/test_pytorch_vgg_network_transform.py
new file mode 100644
index 0000000..1aa6ef5
--- /dev/null
+++ b/test/test_pytorch_vgg_network_transform.py
@@ -0,0 +1,17 @@
+from poet.power_computation import InputLayer
+import torch
+from poet.architectures.network_transform_pytorch import network_transform
+
+# transforms VGG Model network layers into POET computation layers
+
+batch_size = (1,)
+input_shape = (3, 32, 32)
+num_classes = 10
+layers = [InputLayer((batch_size, *input_shape))]
+
+# VGG16 model transformation - https://download.pytorch.org/models/vgg16-397923af.pth
+final_layers = network_transform(
+    torch.hub.load("pytorch/vision:v0.10.0", "vgg16", pretrained=True), layers, batch_size, num_classes, input_shape
+)
+print(final_layers)
+assert len(final_layers) > 0
diff --git a/test/test_tensorflow_resnet_network_transform.py b/test/test_tensorflow_resnet_network_transform.py
new file mode 100644
index 0000000..a3eda61
--- /dev/null
+++ b/test/test_tensorflow_resnet_network_transform.py
@@ -0,0 +1,82 @@
+from poet.architectures.network_transform_tensorflow import network_transform
+from poet.power_computation import InputLayer
+import tensorflow as tf
+
+# transforms ResNet Model network layers into POET computation layers
+
+
+# defining ResNet model in TensorFlow
+def BasicBlock(inputs, num_channels, kernel_size, num_blocks, skip_blocks, name):
+    """Basic residual block"""
+    x = inputs
+    for i in range(num_blocks):
+        if i not in skip_blocks:
+            x1 = ConvNormRelu(x, num_channels, kernel_size, strides=[1, 1], name=name + "." + str(i))
+            x = tf.keras.layers.Add()([x, x1])
+            x = tf.keras.layers.Activation("relu")(x)
+    return x
+
+
+def BasicBlockDown(inputs, num_channels, kernel_size, name):
+    """Residual block with strided downsampling"""
+    x = inputs
+    x1 = ConvNormRelu(x, num_channels, kernel_size, strides=[2, 1], name=name + ".0")
+    x = tf.keras.layers.Conv2D(
+        num_channels, kernel_size=1, strides=2, padding="same", activation="linear", use_bias=False, name=name + ".0.downsample.0"
+    )(x)
+    x = tf.keras.layers.BatchNormalization(momentum=0.1, epsilon=1e-5, name=name + ".0.downsample.1")(x)
+    x = tf.keras.layers.Add()([x, x1])
+    x = tf.keras.layers.Activation("relu")(x)
+    return x
+
+
+def ConvNormRelu(x, num_channels, kernel_size, strides, name):
+    """Layer consisting of 2 consecutive batch normalizations with 1 first relu"""
+    if strides[0] == 2:
+        x = tf.keras.layers.ZeroPadding2D(padding=(1, 1), name=name + ".pad")(x)
+        x = tf.keras.layers.Conv2D(
+            num_channels, kernel_size, strides[0], padding="valid", activation="linear", use_bias=False, name=name + ".conv1"
+        )(x)
+    else:
+        x = tf.keras.layers.Conv2D(
+            num_channels, kernel_size, strides[0], padding="same", activation="linear", use_bias=False, name=name + ".conv1"
+        )(x)
+    x = tf.keras.layers.BatchNormalization(momentum=0.1, epsilon=1e-5, name=name + ".bn1")(x)
+    x = tf.keras.layers.Activation("relu")(x)
+    x = tf.keras.layers.Conv2D(
+        num_channels, kernel_size, strides[1], padding="same", activation="linear", use_bias=False, name=name + ".conv2"
+    )(x)
+    x = tf.keras.layers.BatchNormalization(momentum=0.1, epsilon=1e-5, name=name + ".bn2")(x)
+    return x
+
+
+def ResNet18(inputs):
+    x = tf.keras.layers.ZeroPadding2D(padding=(3, 3), name="pad")(inputs)
+    x = tf.keras.layers.Conv2D(filters=64, kernel_size=7, strides=2, padding="valid", activation="linear", use_bias=False, name="conv1")(x)
+    x = tf.keras.layers.BatchNormalization(momentum=0.1, epsilon=1e-5, name="bn1")(x)
+    x = tf.keras.layers.Activation("relu", name="relu")(x)
+    x = tf.keras.layers.ZeroPadding2D(padding=(1, 1), name="pad1")(x)
+    x = tf.keras.layers.MaxPool2D(pool_size=3, strides=2, padding="valid", name="maxpool")(x)
+    x = BasicBlock(x, num_channels=64, kernel_size=3, num_blocks=2, skip_blocks=[], name="layer1")
+    x = BasicBlockDown(x, num_channels=128, kernel_size=3, name="layer2")
+    x = BasicBlock(x, num_channels=128, kernel_size=3, num_blocks=2, skip_blocks=[0], name="layer2")
+    x = BasicBlockDown(x, num_channels=256, kernel_size=3, name="layer3")
+    x = BasicBlock(x, num_channels=256, kernel_size=3, num_blocks=2, skip_blocks=[0], name="layer3")
+    x = BasicBlockDown(x, num_channels=512, kernel_size=3, name="layer4")
+    x = BasicBlock(x, num_channels=512, kernel_size=3, num_blocks=2, skip_blocks=[0], name="layer4")
+    x = tf.keras.layers.GlobalAveragePooling2D(name="avgpool")(x)
+    x = tf.keras.layers.Dense(units=1000, use_bias=True, activation="linear", name="fc")(x)
+    return x
+
+
+inputs = tf.keras.Input((None, None, 3))
+resnet_tf = ResNet18(inputs)
+model = tf.keras.Model(inputs, resnet_tf)
+batch_size = (1,)
+input_shape = (3, 32, 32)
+num_classes = 10
+layers = [InputLayer((batch_size, *input_shape))]
+
+final_layers = network_transform([layer for layer in model.layers], layers, batch_size, num_classes, input_shape)
+print(final_layers)
+assert len(final_layers) > 0