Create app.py

app.py

@@ -0,0 +1,55 @@
+import json
+import spu.utils.distributed as ppd
+
+from time import time
+from datasets import load_dataset
+from transformers import (
+    AutoConfig,
+    AutoImageProcessor,
+    FlaxResNetForImageClassification,
+)
+
+parser = argparse.ArgumentParser(description='distributed driver.')
+parser.add_argument("-c", "--config", default="3pc.json")
+args = parser.parse_args()
+
+with open(args.config, 'r') as file:
+    conf = json.load(file)
+
+ppd.init(conf["nodes"], conf["devices"])
+
+dataset = load_dataset("huggingface/cats-image")
+image = dataset["test"]["image"][0]
+
+processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
+model = FlaxResNetForImageClassification.from_pretrained("microsoft/resnet-50")
+
+inputs = processor(image, return_tensors="jax")["pixel_values"]
+
+
+def run_on_spu(inputs, model):
+    start = time()
+    inputs = ppd.device("P1")(lambda x: x)(inputs)
+    params = ppd.device("P2")(lambda x: x)(model.params)
+    outputs = ppd.device("SPU")(inference)(inputs, params)
+    outputs = ppd.get(outputs)
+    outputs = outputs['logits']
+    predicted_class_idx = jax.numpy.argmax(outputs, axis=-1)
+    print(f"Elapsed time:{time() - start}")
+    print("Predicted class:", model.config.id2label[predicted_class_idx.item()])
+
+
+def run_on_cpu(inputs, model):
+    start = time()
+    outputs = inference(inputs, model.params)
+    outputs = outputs['logits']
+    predicted_class_idx = jax.numpy.argmax(outputs, axis=-1)
+    print(f"Elapsed time:{time() - start}")
+    print("Predicted class:", model.config.id2label[predicted_class_idx.item()])
+
+
+if __name__ == "__main__":
+    print("Run on CPU\n------\n")
+    run_on_cpu(inputs, model)
+    print("Run on SPU\n------\n")
+    run_on_spu(inputs, model)