Spaces:
Runtime error
Runtime error
File size: 1,715 Bytes
59042ec |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 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) |