Merge pull request #64 from msbaines/onnx

[Python] add initial support for onnx
facebookresearch · Nov 17, 2020 · 1d27117 · 1d27117
2 parents 6b7a94b + bb67425
commit 1d27117
Show file tree

Hide file tree

Showing 11 changed files with 1,200 additions and 0 deletions.
diff --git a/examples/python/onnx/accuracy.py b/examples/python/onnx/accuracy.py
@@ -0,0 +1,23 @@
+# Copyright 2020 Stanford University, Los Alamos National Laboratory
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from enum import Enum
+
+class ModelAccuracy(Enum):
+  MNIST_MLP = 90
+  MNIST_CNN = 90
+  REUTERS_MLP = 90
+  CIFAR10_CNN = 90
+  CIFAR10_ALEXNET = 90
diff --git a/examples/python/onnx/alexnet.py b/examples/python/onnx/alexnet.py
@@ -0,0 +1,174 @@
+from flexflow.core import *
+from flexflow.keras.datasets import cifar10
+from flexflow.onnx.model import ONNXModel
+
+from accuracy import ModelAccuracy
+from PIL import Image
+
+def top_level_task():
+  ffconfig = FFConfig()
+  alexnetconfig = NetConfig()
+  print(alexnetconfig.dataset_path)
+  ffconfig.parse_args()
+  print("Python API batchSize(%d) workersPerNodes(%d) numNodes(%d)" %(ffconfig.get_batch_size(), ffconfig.get_workers_per_node(), ffconfig.get_num_nodes()))
+  ffmodel = FFModel(ffconfig)
+
+  dims_input = [ffconfig.get_batch_size(), 3, 229, 229]
+  #print(dims)
+  input = ffmodel.create_tensor(dims_input, "", DataType.DT_FLOAT)
+
+  # dims_label = [ffconfig.get_batch_size(), 1]
+  # #print(dims)
+  # label = ffmodel.create_tensor(dims_label, "", DataType.DT_INT32)
+
+  onnx_model = ONNXModel("alexnet.onnx")
+  t = onnx_model.apply(ffmodel, {"input.1": input})
+
+  ffoptimizer = SGDOptimizer(ffmodel, 0.01)
+  ffmodel.set_sgd_optimizer(ffoptimizer)
+  ffmodel.compile(loss_type=LossType.LOSS_SPARSE_CATEGORICAL_CROSSENTROPY, metrics=[MetricsType.METRICS_ACCURACY, MetricsType.METRICS_SPARSE_CATEGORICAL_CROSSENTROPY])
+  label = ffmodel.get_label_tensor()
+
+  use_external = True
+  if (use_external == True):
+    num_samples = 10000
+
+    (x_train, y_train), (x_test, y_test) = cifar10.load_data(num_samples)
+
+    full_input_np = np.zeros((num_samples, 3, 229, 229), dtype=np.float32)
+
+    for i in range(0, num_samples):
+      image = x_train[i, :, :, :]
+      image = image.transpose(1, 2, 0)
+      pil_image = Image.fromarray(image)
+      pil_image = pil_image.resize((229,229), Image.NEAREST)
+      image = np.array(pil_image, dtype=np.float32)
+      image = image.transpose(2, 0, 1)
+      full_input_np[i, :, :, :] = image
+      if (i == 0):
+        print(image)
+
+
+    full_input_np /= 255
+    print(full_input_np.shape)
+    print(full_input_np.__array_interface__["strides"])
+    print(full_input_np[0,:, :, :])
+
+    y_train = y_train.astype('int32')
+    full_label_np = y_train
+
+    dims_full_input = [num_samples, 3, 229, 229]
+    full_input = ffmodel.create_tensor(dims_full_input, "", DataType.DT_FLOAT)
+
+    dims_full_label = [num_samples, 1]
+    full_label = ffmodel.create_tensor(dims_full_label, "", DataType.DT_INT32)
+
+    full_input.attach_numpy_array(ffconfig, full_input_np)
+    full_label.attach_numpy_array(ffconfig, full_label_np)
+
+    dataloader_input = SingleDataLoader(ffmodel, input, full_input, num_samples, DataType.DT_FLOAT)
+    dataloader_label = SingleDataLoader(ffmodel, label, full_label, num_samples, DataType.DT_INT32)
+    #dataloader = DataLoader4D(ffmodel, input, label, full_input, full_label, num_samples)
+
+    full_input.detach_numpy_array(ffconfig)
+    full_label.detach_numpy_array(ffconfig)
+
+    num_samples = dataloader_input.get_num_samples()
+    assert dataloader_input.get_num_samples() == dataloader_label.get_num_samples()
+
+  else:
+    # Data Loader
+    dataloader = DataLoader4D(ffmodel, input, label, ffnetconfig=alexnetconfig)
+    num_samples = dataloader.get_num_samples()
+
+  # input.inline_map(ffconfig)
+  # input_array = input.get_array(ffconfig, DataType.DT_FLOAT)
+  # input_array *= 1.0
+  # print(input_array.shape)
+  # input.inline_unmap(ffconfig)
+  # label.inline_map(ffconfig)
+  # label.inline_unmap(ffconfig)
+
+  ffmodel.init_layers()
+
+ #  conv_2d1 = ffmodel.get_layer_by_id(11)
+ #  cbias_tensor = conv_2d1.get_weight_tensor()
+ #  input_tensor = conv_2d1.get_input_tensor_by_id(0)
+ #  cbias_tensor.inline_map(ffconfig)
+ #  cbias = cbias_tensor.get_array(ffconfig, DataType.DT_FLOAT)
+ # # cbias += 0.125
+ #  print(cbias.shape)
+ #  #print(cbias)
+ #  cbias_tensor.inline_unmap(ffconfig)
+
+  #use_external = False
+
+  epochs = ffconfig.get_epochs()
+
+  ts_start = ffconfig.get_current_time()
+  for epoch in range(0,epochs):
+    if (use_external == True):
+      dataloader_input.reset()
+      dataloader_label.reset()
+    else:
+      dataloader.reset()
+    ffmodel.reset_metrics()
+    iterations = int(num_samples / ffconfig.get_batch_size())
+    print(iterations, num_samples)
+
+    for iter in range(0, int(iterations)):
+      if (use_external == True):
+        dataloader_input.next_batch(ffmodel)
+        dataloader_label.next_batch(ffmodel)
+      else:
+        dataloader.next_batch(ffmodel)
+      if (epoch > 0):
+        ffconfig.begin_trace(111)
+      ffmodel.forward()
+      ffmodel.zero_gradients()
+      ffmodel.backward()
+      ffmodel.update()
+      if (epoch > 0):
+        ffconfig.end_trace(111)
+
+  ts_end = ffconfig.get_current_time()
+  run_time = 1e-6 * (ts_end - ts_start);
+  print("epochs %d, ELAPSED TIME = %.4fs, THROUGHPUT = %.2f samples/s\n" %(epochs, run_time, num_samples * epochs / run_time));
+  perf_metrics = ffmodel.get_perf_metrics()
+  accuracy = perf_metrics.get_accuracy()
+  if accuracy < ModelAccuracy.CIFAR10_ALEXNET.value:
+    assert 0, 'Check Accuracy'
+
+  #ffmodel.print_layers(13)
+
+  conv_2d1 = ffmodel.get_layer_by_id(0)
+  #cbias_tensor = conv_2d1.get_input_tensor()
+  cbias_tensor = conv_2d1.get_input_tensor()
+  cbias_tensor.inline_map(ffconfig)
+  cbias = cbias_tensor.get_flat_array(ffconfig, DataType.DT_FLOAT)
+  print(cbias.shape)
+  print(cbias)
+  #save_image(cbias, 2)
+  cbias_tensor.inline_unmap(ffconfig)
+
+  label.inline_map(ffconfig)
+  label_array = label.get_flat_array(ffconfig, DataType.DT_INT32)
+  print(label_array.shape)
+  # print(cbias)
+  print(label_array)
+  label.inline_unmap(ffconfig)
+
+  #ffmodel.print_layers(0)
+
+def save_image(batch_image_array, id):
+  image_array = batch_image_array[id, :, :, :]
+  image_array = image_array.transpose(1, 2, 0)
+  image_array = image_array*255
+  image_array = image_array.astype('uint8')
+  pil_image = Image.fromarray(image_array).convert('RGB')
+  pil_image.save("img.jpeg")
+
+
+if __name__ == "__main__":
+  print("alexnet")
+  top_level_task()
diff --git a/examples/python/onnx/alexnet_pt.py b/examples/python/onnx/alexnet_pt.py
@@ -0,0 +1,41 @@
+import onnx
+import torch
+import torch.nn as nn
+from torch.onnx import TrainingMode
+
+class AlexNet(nn.Module):
+    def __init__(self, num_classes: int = 1000) -> None:
+        super(AlexNet, self).__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            nn.Conv2d(64, 192, kernel_size=5, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            nn.Conv2d(192, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+        )
+        self.classifier = nn.Sequential(
+            nn.Linear(256 * 6 * 6, 4096),
+            nn.ReLU(inplace=True),
+            nn.Linear(4096, 4096),
+            nn.ReLU(inplace=True),
+            nn.Linear(4096, num_classes),
+            nn.Softmax(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+        return x
+
+input = torch.randn(64, 3, 224, 224)
+model = AlexNet()
+torch.onnx.export(model, (input), "alexnet.onnx", export_params=False, training=TrainingMode.TRAINING)
diff --git a/examples/python/onnx/cifar10_cnn.py b/examples/python/onnx/cifar10_cnn.py
@@ -0,0 +1,134 @@
+from flexflow.core import *
+from flexflow.keras.datasets import cifar10
+from flexflow.onnx.model import ONNXModel
+
+from accuracy import ModelAccuracy
+
+def top_level_task():
+  ffconfig = FFConfig()
+  alexnetconfig = NetConfig()
+  print(alexnetconfig.dataset_path)
+  ffconfig.parse_args()
+  print("Python API batchSize(%d) workersPerNodes(%d) numNodes(%d)" %(ffconfig.get_batch_size(), ffconfig.get_workers_per_node(), ffconfig.get_num_nodes()))
+  ffmodel = FFModel(ffconfig)
+
+  dims_input = [ffconfig.get_batch_size(), 3, 32, 32]
+  input = ffmodel.create_tensor(dims_input, "", DataType.DT_FLOAT)
+
+  # dims_label = [ffconfig.get_batch_size(), 1]
+  # label = ffmodel.create_tensor(dims_label, "", DataType.DT_INT32)
+
+  onnx_model = ONNXModel("cifar10_cnn.onnx")
+  t = onnx_model.apply(ffmodel, {"input.1": input})
+
+  ffoptimizer = SGDOptimizer(ffmodel, 0.01)
+  ffmodel.set_sgd_optimizer(ffoptimizer)
+  ffmodel.compile(loss_type=LossType.LOSS_SPARSE_CATEGORICAL_CROSSENTROPY, metrics=[MetricsType.METRICS_ACCURACY, MetricsType.METRICS_SPARSE_CATEGORICAL_CROSSENTROPY])
+  label = ffmodel.get_label_tensor()
+
+  use_external = True
+  if (use_external == True):
+    num_samples = 10000
+
+    (x_train, y_train), (x_test, y_test) = cifar10.load_data(num_samples)
+
+    x_train = x_train.astype('float32')
+    x_train /= 255
+    full_input_array = x_train
+    print(full_input_array.__array_interface__["strides"])
+
+    y_train = y_train.astype('int32')
+    full_label_array = y_train
+
+    print(full_input_array.__array_interface__["strides"])
+    print(full_input_array.shape, full_label_array.shape)
+    #print(full_input_array[0,:,:,:])
+    #print(full_label_array[0, 0:64])
+    print(full_label_array.__array_interface__["strides"])
+
+    dims_full_input = [num_samples, 3, 32, 32]
+    full_input = ffmodel.create_tensor(dims_full_input, "", DataType.DT_FLOAT)
+
+    dims_full_label = [num_samples, 1]
+    full_label = ffmodel.create_tensor(dims_full_label, "", DataType.DT_INT32)
+
+    full_input.attach_numpy_array(ffconfig, full_input_array)
+    full_label.attach_numpy_array(ffconfig, full_label_array)
+
+    dataloader_input = SingleDataLoader(ffmodel, input, full_input, num_samples, DataType.DT_FLOAT)
+    dataloader_label = SingleDataLoader(ffmodel, label, full_label, num_samples, DataType.DT_INT32)
+
+    full_input.detach_numpy_array(ffconfig)
+    full_label.detach_numpy_array(ffconfig)
+
+    num_samples = dataloader_input.get_num_samples()
+  else:
+    # Data Loader
+    dataloader = DataLoader4D(ffmodel, input, label, ffnetconfig=alexnetconfig)
+    num_samples = dataloader.get_num_samples()
+
+  ffmodel.init_layers()
+
+  epochs = ffconfig.get_epochs()
+  #epochs = 10
+
+  ts_start = ffconfig.get_current_time()
+  for epoch in range(0,epochs):
+    if (use_external == True):
+      dataloader_input.reset()
+      dataloader_label.reset()
+    else:
+      dataloader.reset()
+    ffmodel.reset_metrics()
+    iterations = int(num_samples / ffconfig.get_batch_size())
+    print(iterations, num_samples)
+
+    for iter in range(0, int(iterations)):
+      # if (len(alexnetconfig.dataset_path) == 0):
+      #   if (iter == 0 and epoch == 0):
+      #     dataloader.next_batch(ffmodel)
+      # else:
+    #  dataloader.next_batch(ffmodel)
+      if (use_external == True):
+        dataloader_input.next_batch(ffmodel)
+        dataloader_label.next_batch(ffmodel)
+      else:
+        dataloader.next_batch(ffmodel)
+      if (epoch > 0):
+        ffconfig.begin_trace(111)
+      ffmodel.forward()
+      ffmodel.zero_gradients()
+      ffmodel.backward()
+      ffmodel.update()
+      if (epoch > 0):
+        ffconfig.end_trace(111)
+
+  ts_end = ffconfig.get_current_time()
+  run_time = 1e-6 * (ts_end - ts_start);
+  print("epochs %d, ELAPSED TIME = %.4fs, THROUGHPUT = %.2f samples/s\n" %(epochs, run_time, num_samples * epochs / run_time));
+
+  perf_metrics = ffmodel.get_perf_metrics()
+  accuracy = perf_metrics.get_accuracy()
+  if accuracy < ModelAccuracy.CIFAR10_CNN.value:
+    assert 0, 'Check Accuracy'
+
+  conv_2d1 = ffmodel.get_layer_by_id(0)
+  cbias_tensor = conv_2d1.get_input_tensor()
+  #cbias_tensor = conv_2d1.get_output_tensor()
+  cbias_tensor.inline_map(ffconfig)
+  cbias = cbias_tensor.get_flat_array(ffconfig, DataType.DT_FLOAT)
+  print(cbias.shape)
+  print(cbias)
+  cbias_tensor.inline_unmap(ffconfig)
+
+  label.inline_map(ffconfig)
+  label_array = label.get_flat_array(ffconfig, DataType.DT_INT32)
+  print(label_array.shape)
+  # print(cbias)
+  print(label_array)
+  label.inline_unmap(ffconfig)
+
+
+if __name__ == "__main__":
+  print("cifar10 cnn")
+  top_level_task()