Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,8 @@ openvino-mobilenet-raw/mobilenet.bin
openvino-mobilenet-raw/mobilenet.xml
openvino-mobilenet-raw/tensor-1x224x224x3-f32.bgr

pytorch-mobilenet-image/fixtures/input.jpg
# pytorch-mobilenet-image/mobilenet.pt

.DS_Store
Cargo.lock
67 changes: 67 additions & 0 deletions pytorch-mobilenet-image/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
# Mobilenet example for WASI-NN

This package is a high-level Rust bindings for [wasi-nn] example of Mobilenet.

[wasi-nn]: https://github.com/WebAssembly/wasi-nn

## Dependencies

This crate depends on the `wasi-nn` in the `Cargo.toml`:

```toml
[dependencies]
wasi-nn = "0.1.0"
```

## Build

Compile the application to WebAssembly:

```bash
cargo build --target=wasm32-wasi --release
```

The output WASM file will be at `target/wasm32-wasi/release/wasmedge-wasinn-example-mobilenet-image.wasm`.
To speed up the image processing, we can enable the AOT mode in WasmEdge with:

```bash
wasmedgec rust/target/wasm32-wasi/release/wasmedge-wasinn-example-mobilenet-image.wasm wasmedge-wasinn-example-mobilenet-image.wasm
```

## Run

First generate the fixture of the pre-trained mobilenet with the script:

```bash
./download_data.sh fixtures && cd fixtures
python -m pip install -r requirements.txt
# generate the model fixture
python generate_mobilenet.py
```

(or you can use the pre-generated fixture in `fixtures/mobilenet.pt`)

The above will download a testing image `input.jpg`
![](https://github.com/bytecodealliance/wasi-nn/raw/main/rust/images/1.jpg)
as well as a pre-trained mobilenet model, then convert the model into the torchscript model for C++.

And execute the WASM with the `wasmedge` with PyTorch supporting:

```bash
wasmedge --dir .:. wasmedge-wasinn-example-mobilenet-image.wasm fixtures/mobilenet.pt input.jpg
```

You will get the output:

```console
Read torchscript binaries, size in bytes: 14376924
Loaded graph into wasi-nn with ID: 0
Created wasi-nn execution context with ID: 0
Read input tensor, size in bytes: 602112
Executed graph inference
1.) [954](20.6681)banana
2.) [940](12.1483)spaghetti squash
3.) [951](11.5748)lemon
4.) [950](10.4899)orange
5.) [953](9.4834)pineapple, ananas
```
6 changes: 6 additions & 0 deletions pytorch-mobilenet-image/download_data.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
FIXTURE=https://github.com/intel/openvino-rs/raw/v0.3.3/crates/openvino/tests/fixtures/mobilenet
TODIR=$1

if [ ! -f $TODIR/input.jpg ]; then
wget --no-clobber https://github.com/bytecodealliance/wasi-nn/raw/main/rust/images/1.jpg -O $TODIR/input.jpg
fi
Empty file.
57 changes: 57 additions & 0 deletions pytorch-mobilenet-image/fixtures/generate_mobilenet.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@

# adapted from https://pytorch.org/hub/pytorch_vision_mobilenet_v2/

import torch
import struct
# Download an example image from the pytorch website
url, filename = (
"https://github.com/bytecodealliance/wasi-nn/raw/main/rust/images/1.jpg", "../input.jpg")
# import urllib
# try:
# urllib.URLopener().retrieve(url, filename)
# except:
# urllib.request.urlretrieve(url, filename)

# sample execution (requires torchvision)
model = torch.hub.load('pytorch/vision:v0.10.0',
'mobilenet_v2', pretrained=True)
model.eval()

from PIL import Image
from torchvision import transforms
input_image = Image.open(filename)
print(input_image.mode)
preprocess = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
# create a mini-batch as expected by the model
input_batch = input_tensor.unsqueeze(0)
with open("image-1-3-244-244.rgb", 'wb') as f:
order_data = input_batch.reshape(-1)
for d in order_data:
d = d.item()
f.write(struct.pack('f', d))


# move the input and model to GPU for speed if available
if torch.cuda.is_available():
input_batch = input_batch.to('cuda')
model.to('cuda')

with torch.no_grad():
output = model(input_batch)
# Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes
# The output has unnormalized scores. To get probabilities, you can run a softmax on it.
probabilities = torch.nn.functional.softmax(output[0], dim=0)
# print(probabilities)

with open("imagenet_classes.txt", "r") as f:
categories = [s.strip() for s in f.readlines()]
# Show top categories per image
top5_prob, top5_catid = torch.topk(probabilities, 5)
for i in range(top5_prob.size(0)):
print(top5_catid[i], categories[top5_catid[i]], top5_prob[i].item())
Binary file not shown.
Loading