From c46a0084d5140e3bcaf7140d0ed89cdf1bbcbc19 Mon Sep 17 00:00:00 2001 From: Suyash20012000 Date: Sun, 9 Feb 2025 08:29:59 -0500 Subject: [PATCH 1/2] Adding support for accepting pytorch model config as input. --- scalesim/scale.py | 28 ++++++++++ scalesim/torch_to_topo.py | 108 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+) create mode 100644 scalesim/torch_to_topo.py diff --git a/scalesim/scale.py b/scalesim/scale.py index 1aac3c25a..2bb9b7feb 100644 --- a/scalesim/scale.py +++ b/scalesim/scale.py @@ -1,6 +1,12 @@ import argparse +import os + +import torch +import torch.nn as nn from scalesim.scale_sim import scalesim +from torch_to_topo import create_conv_topo +from torchvision import models, datasets if __name__ == '__main__': parser = argparse.ArgumentParser() @@ -20,6 +26,15 @@ default="conv", help="Type of input topology, gemm: MNK, conv: conv" ) + # Changes from torch_to_topo files + parser.add_argument('-m', metavar='PyTorch module name', type=str, + default='', + help="Class name of PyTorch module to turn into SCALESim topology" + ) + parser.add_argument('-l', metavar='Include linear layers?', type=bool, + default=False, + help="Boolean deciding whether or not to include nn.Linear layers in topology" + ) args = parser.parse_args() topology = args.t @@ -27,6 +42,19 @@ logpath = args.p inp_type = args.i + if args.m != '': + # Changes from torch_to_topo files + path = '../topologies/custom/' + if not os.path.exists(path): + os.makedirs(path) + include_lin = args.l + model_name = args.m + model = getattr(models, model_name)(pretrained=True) + file_name = model_name + '.csv' + create_conv_topo(model, path=path, filename=file_name, include_lin=include_lin) + + topology = path + file_name + gemm_input = False if inp_type == 'gemm': gemm_input = True diff --git a/scalesim/torch_to_topo.py b/scalesim/torch_to_topo.py new file mode 100644 index 000000000..2fcc73a1f --- /dev/null +++ b/scalesim/torch_to_topo.py @@ -0,0 +1,108 @@ + +import argparse +import os + +import torch +import torch.nn as nn + +from torchvision import models, datasets +from torchvision.transforms import ToTensor +from torchinfo import summary + +from scalesim import topology_utils + + +base_dir = '../topologies/custom' + + +# creates topology file from PyTorch module +def create_conv_topo(model, path=base_dir, filename='', input_dim=(2, 3, 224, 224), include_lin=False): + if filename == '': + filename = model.__class__.__name__+'_custom.csv' + + topo = topology_utils.topologies() + prev_arr, curr_arr = [], [] + last = 0 + + model_stats = summary(model, input_size=input_dim, col_names=('input_size','output_size','num_params','kernel_size'), verbose=1) + idx = 1 + for layer in model_stats.summary_list: + if layer.class_name == 'Conv2d': + if prev_arr: + prev_arr[6] = last + topo.topo_arrays.append(prev_arr) + idx += 1 + curr_arr.append(layer.class_name+'_{}'.format(idx)) # layer name + curr_arr.append(layer.input_size[2]) # IFMAP height + curr_arr.append(layer.input_size[3]) # IFMAP width + curr_arr.append(layer.kernel_size[0]) # filter height + curr_arr.append(layer.kernel_size[1]) # filter width + curr_arr.append(layer.input_size[1]) # channels + curr_arr.append(0) # output channels (num filters) + curr_arr.append(layer.module.stride[0]) # stride + prev_arr = curr_arr + curr_arr = [] + elif layer.class_name == 'Linear' and include_lin: + if prev_arr: + prev_arr[6] = last + topo.topo_arrays.append(prev_arr) + idx += 1 + curr_arr.append(layer.class_name+'_{}'.format(idx)) # layer name + print(layer.input_size) + curr_arr.append(1) # IFMAP height + curr_arr.append(1) # IFMAP width + curr_arr.append(1) # filter height + curr_arr.append(1) # filter width + curr_arr.append(layer.input_size[1]) # channels + curr_arr.append(0) # output channels (num filters) + curr_arr.append(1) # stride + prev_arr = curr_arr + curr_arr = [] + last = layer.input_size[1] + # last layer + if prev_arr: + prev_arr[6] = last + topo.topo_arrays.append(prev_arr) + + topo.topo_load_flag = True + topo.write_topo_file(path=path, filename=filename) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('-m', metavar='PyTorch module name', type=str, + default='alexnet', + help="Class name of PyTorch module to turn into SCALESim topology" + ) + parser.add_argument('-p', metavar='Path to directory', type=str, + default='../topologies/custom', + help="Directory at which to save the created topology file" + ) + parser.add_argument('-f', metavar='File name', type=str, + default='alexnet_custom.csv', + help="File name for SCALESim topology" + ) + parser.add_argument('-l', metavar='Include linear layers?', type=bool, + default=False, + help="Boolean deciding whether or not to include nn.Linear layers in topology" + ) + + + args = parser.parse_args() + if not os.path.exists(base_dir): + os.mkdir(base_dir) + + model_name = args.m + print(model_name) + print(args.p) + try: + model = getattr(models, model_name)(pretrained=True) + except: + print('PyTorch module does not exist') + exit() + + path = args.p + file_name = args.f + include_lin = args.l + + create_conv_topo(model, path=path, filename=file_name, include_lin=include_lin) \ No newline at end of file From c71d438918de81c24999dc96bddf27de73bfd27a Mon Sep 17 00:00:00 2001 From: Suyash20012000 Date: Thu, 13 Feb 2025 19:03:50 -0500 Subject: [PATCH 2/2] Updating requirements.txt to add pytorch requirements. --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 80a346c8b..f422b9ee2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ numpy pandas cython setuptools +torch \ No newline at end of file