+Ivy AlexNet demo#
+In this demo, we show how an AlexNet model written in Ivy native code, can be used for image classification, and integrated with all three of the major ML frameworks: PyTorch, TensorFlow and JAX.
+Installation#
+Since we want the packages to be available after installing, after running the first cell, the notebook will automatically restart.
+You can then do Runtime -> Run all after the notebook has restarted, to run all of the cells.
+Make sure you run this demo with GPU enabled!
+
+
+
+
+
+
+[1]:
+!pip install -q ivy
+!pip install -q dm-haiku
+!git clone https://github.com/unifyai/models.git --depth 1
+
+# Installing models package from cloned repository! 😄
+!cd models/ && pip install .
+
+
+
+
+
+
+
+WARNING: Ignoring invalid distribution -ocutils (/opt/homebrew/lib/python3.10/site-packages)
+WARNING: Ignoring invalid distribution -ocutils (/opt/homebrew/lib/python3.10/site-packages)
+WARNING: Ignoring invalid distribution -ocutils (/opt/homebrew/lib/python3.10/site-packages)
+WARNING: Ignoring invalid distribution -ocutils (/opt/homebrew/lib/python3.10/site-packages)
+ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
+tensorflow-macos 2.15.0 requires ml-dtypes~=0.2.0, but you have ml-dtypes 0.4.0 which is incompatible.
+Cloning into 'models'...
+remote: Enumerating objects: 192, done.
+remote: Counting objects: 100% (192/192), done.
+remote: Compressing objects: 100% (156/156), done.
+remote: Total 192 (delta 37), reused 104 (delta 13), pack-reused 0
+Receiving objects: 100% (192/192), 1.83 MiB | 3.45 MiB/s, done.
+Resolving deltas: 100% (37/37), done.
+WARNING: Ignoring invalid distribution -ocutils (/opt/homebrew/lib/python3.10/site-packages)
+Processing /Users/samuelarmstrong/Documents/ivy/demos/examples_and_demos_cpu/models
+ Preparing metadata (setup.py) ... done
+Requirement already satisfied: ivy in /opt/homebrew/lib/python3.10/site-packages (from ivy-models==1.1.10) (0.0.8.0)
+Requirement already satisfied: scipy in /opt/homebrew/lib/python3.10/site-packages (from ivy-models==1.1.10) (1.12.0)
+Requirement already satisfied: numpy in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (1.26.4)
+Requirement already satisfied: einops in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (0.7.0)
+Requirement already satisfied: psutil in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (5.9.8)
+Requirement already satisfied: termcolor in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (2.4.0)
+Requirement already satisfied: colorama in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (0.4.6)
+Requirement already satisfied: packaging in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (24.0)
+Requirement already satisfied: nvidia-ml-py in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (12.535.133)
+Requirement already satisfied: diskcache in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (5.6.3)
+Requirement already satisfied: google-auth in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (2.29.0)
+Requirement already satisfied: urllib3<2.0 in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (1.26.18)
+Requirement already satisfied: requests in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (2.31.0)
+Requirement already satisfied: pyvis in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (0.3.2)
+Requirement already satisfied: dill in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (0.3.8)
+Requirement already satisfied: astunparse in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (1.6.3)
+Requirement already satisfied: ml-dtypes in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (0.4.0)
+Requirement already satisfied: cloudpickle in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (3.0.0)
+Requirement already satisfied: gast in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (0.5.4)
+Requirement already satisfied: tqdm in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (4.66.2)
+Requirement already satisfied: cryptography in /opt/homebrew/lib/python3.10/site-packages (from ivy->ivy-models==1.1.10) (42.0.5)
+Requirement already satisfied: wheel<1.0,>=0.23.0 in /opt/homebrew/lib/python3.10/site-packages (from astunparse->ivy->ivy-models==1.1.10) (0.43.0)
+Requirement already satisfied: six<2.0,>=1.6.1 in /opt/homebrew/lib/python3.10/site-packages (from astunparse->ivy->ivy-models==1.1.10) (1.16.0)
+Requirement already satisfied: cffi>=1.12 in /opt/homebrew/lib/python3.10/site-packages (from cryptography->ivy->ivy-models==1.1.10) (1.16.0)
+Requirement already satisfied: cachetools<6.0,>=2.0.0 in /opt/homebrew/lib/python3.10/site-packages (from google-auth->ivy->ivy-models==1.1.10) (5.3.3)
+Requirement already satisfied: pyasn1-modules>=0.2.1 in /opt/homebrew/lib/python3.10/site-packages (from google-auth->ivy->ivy-models==1.1.10) (0.3.0)
+Requirement already satisfied: rsa<5,>=3.1.4 in /opt/homebrew/lib/python3.10/site-packages (from google-auth->ivy->ivy-models==1.1.10) (4.9)
+Requirement already satisfied: ipython>=5.3.0 in /opt/homebrew/lib/python3.10/site-packages (from pyvis->ivy->ivy-models==1.1.10) (8.22.2)
+Requirement already satisfied: jinja2>=2.9.6 in /opt/homebrew/lib/python3.10/site-packages (from pyvis->ivy->ivy-models==1.1.10) (3.1.3)
+Requirement already satisfied: jsonpickle>=1.4.1 in /opt/homebrew/lib/python3.10/site-packages (from pyvis->ivy->ivy-models==1.1.10) (3.0.3)
+Requirement already satisfied: networkx>=1.11 in /opt/homebrew/lib/python3.10/site-packages (from pyvis->ivy->ivy-models==1.1.10) (3.2.1)
+Requirement already satisfied: charset-normalizer<4,>=2 in /opt/homebrew/lib/python3.10/site-packages (from requests->ivy->ivy-models==1.1.10) (3.3.2)
+Requirement already satisfied: idna<4,>=2.5 in /opt/homebrew/lib/python3.10/site-packages (from requests->ivy->ivy-models==1.1.10) (3.6)
+Requirement already satisfied: certifi>=2017.4.17 in /opt/homebrew/lib/python3.10/site-packages (from requests->ivy->ivy-models==1.1.10) (2024.2.2)
+Requirement already satisfied: pycparser in /opt/homebrew/lib/python3.10/site-packages (from cffi>=1.12->cryptography->ivy->ivy-models==1.1.10) (2.21)
+Requirement already satisfied: decorator in /opt/homebrew/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (5.1.1)
+Requirement already satisfied: jedi>=0.16 in /opt/homebrew/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (0.19.1)
+Requirement already satisfied: matplotlib-inline in /opt/homebrew/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (0.1.6)
+Requirement already satisfied: prompt-toolkit<3.1.0,>=3.0.41 in /opt/homebrew/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (3.0.43)
+Requirement already satisfied: pygments>=2.4.0 in /opt/homebrew/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (2.17.2)
+Requirement already satisfied: stack-data in /opt/homebrew/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (0.6.3)
+Requirement already satisfied: traitlets>=5.13.0 in /opt/homebrew/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (5.14.2)
+Requirement already satisfied: exceptiongroup in /opt/homebrew/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (1.2.0)
+Requirement already satisfied: pexpect>4.3 in /opt/homebrew/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (4.9.0)
+Requirement already satisfied: MarkupSafe>=2.0 in /opt/homebrew/lib/python3.10/site-packages (from jinja2>=2.9.6->pyvis->ivy->ivy-models==1.1.10) (2.1.5)
+Requirement already satisfied: pyasn1<0.6.0,>=0.4.6 in /opt/homebrew/lib/python3.10/site-packages (from pyasn1-modules>=0.2.1->google-auth->ivy->ivy-models==1.1.10) (0.5.1)
+Requirement already satisfied: parso<0.9.0,>=0.8.3 in /opt/homebrew/lib/python3.10/site-packages (from jedi>=0.16->ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (0.8.3)
+Requirement already satisfied: ptyprocess>=0.5 in /opt/homebrew/lib/python3.10/site-packages (from pexpect>4.3->ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (0.7.0)
+Requirement already satisfied: wcwidth in /opt/homebrew/lib/python3.10/site-packages (from prompt-toolkit<3.1.0,>=3.0.41->ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (0.2.13)
+Requirement already satisfied: executing>=1.2.0 in /opt/homebrew/lib/python3.10/site-packages (from stack-data->ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (2.0.1)
+Requirement already satisfied: asttokens>=2.1.0 in /opt/homebrew/lib/python3.10/site-packages (from stack-data->ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (2.4.1)
+Requirement already satisfied: pure-eval in /opt/homebrew/lib/python3.10/site-packages (from stack-data->ipython>=5.3.0->pyvis->ivy->ivy-models==1.1.10) (0.2.2)
+Building wheels for collected packages: ivy-models
+ Building wheel for ivy-models (setup.py) ... done
+ Created wheel for ivy-models: filename=ivy_models-1.1.10-py3-none-any.whl size=76449 sha256=4ef86060439480c8cdd692e30d269e68540c3728a30c4a372981b0c5c0cbc214
+ Stored in directory: /private/var/folders/3x/7zt1qbl12mn7zq12fzzv6xh80000gn/T/pip-ephem-wheel-cache-abb7vdwj/wheels/01/2d/88/adc983ab61e1210a8d2ee2a20d1fc3d7c3e082fcdeabe25595
+Successfully built ivy-models
+WARNING: Ignoring invalid distribution -ocutils (/opt/homebrew/lib/python3.10/site-packages)
+Installing collected packages: ivy-models
+Successfully installed ivy-models-1.1.10
+Data Preparation#
+First we need to download the ImageNet classes and preprocess the image.
+
+
+
+
+
+
+[2]:
+!wget https://github.com/raw/pytorch/hub/master/imagenet_classes.txt
+with open("imagenet_classes.txt", "r") as f:
+ categories = [s.strip() for s in f.readlines()]
+
+
+
+
+
+
+
+zsh:1: command not found: wget
+
+
+
+
+
+
+
+---------------------------------------------------------------------------
+FileNotFoundError Traceback (most recent call last)
+Cell In[2], line 2
+ 1 get_ipython().system('wget https://github.com/raw/pytorch/hub/master/imagenet_classes.txt')
+----> 2 with open("imagenet_classes.txt", "r") as f:
+ 3 categories = [s.strip() for s in f.readlines()]
+
+File /opt/homebrew/lib/python3.10/site-packages/IPython/core/interactiveshell.py:324, in _modified_open(file, *args, **kwargs)
+ 317 if file in {0, 1, 2}:
+ 318 raise ValueError(
+ 319 f"IPython won't let you open fd={file} by default "
+ 320 "as it is likely to crash IPython. If you know what you are doing, "
+ 321 "you can use builtins' open."
+ 322 )
+--> 324 return io_open(file, *args, **kwargs)
+
+FileNotFoundError: [Errno 2] No such file or directory: 'imagenet_classes.txt'
+
+
+
+
+
+
+
+The Kernel crashed while executing code in the current cell or a previous cell.
+
+Please review the code in the cell(s) to identify a possible cause of the failure.
+
+Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info.
+
+View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details.
+
+
+
+
+
+
+[ ]:
+!wget https://github.com/raw/unifyai/models/master/images/cat.jpg
+filename = "cat.jpg"
+
+
+
+
+
+
+[ ]:
+# Preprocess torch image
+import jax
+jax.devices()
+import torch
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+import warnings
+import time
+warnings.filterwarnings('ignore')
+
+preprocess = transforms.Compose([
+ transforms.Resize(256),
+ transforms.CenterCrop(224),
+ transforms.ToTensor(),
+ transforms.Normalize(
+ mean=[0.485, 0.456, 0.406],
+ std=[0.229, 0.224, 0.225]
+)])
+torch_img = Image.open(filename)
+torch_img = preprocess(torch_img)
+torch_img = torch.unsqueeze(torch_img, 0)
+
+img = torch_img.numpy()
+
+
+
+
+
+
+[ ]:
+from IPython.display import Image, display
+display(Image(filename))
+
+
+
+
+
+
+
+
+Ivy AlexNet inference in Torch#
+We import the Ivy native implementation of AlexNet. The code for this model is given at the end of this notebook, click here to see it.
+
+
+
+
+
+
+[ ]:
+import ivy
+ivy.set_soft_device_mode(True)
+ivy.set_backend("torch")
+
+from ivy_models.alexnet import alexnet
+ivy_alexnet = alexnet()
+In order to make sure the model is as quick as possible, we can call ivy.trace_graph(). This can take a moment, but is a one-time cost.
+
+
+
+
+
+[ ]:
+ivy_alexnet = ivy.trace_graph(ivy_alexnet, args=(ivy.asarray(torch_img),))
+
+
+
+
+
+
+[ ]:
+output = ivy.softmax(ivy_alexnet(ivy.asarray(img))) # pass the image to the model
+classes = ivy.argsort(output[0], descending=True)[:3] # get the top 3 classes
+logits = ivy.gather(output[0], classes) # get the logits
+
+print("Indices of the top 3 classes are:", classes)
+print("Logits of the top 3 classes are:", logits)
+print("Categories of the top 3 classes are:", [categories[i] for i in classes.to_list()])
+
+
+
+
+
+
+
+Indices of the top 3 classes are: ivy.array([282, 281, 285], dev=gpu:0)
+Logits of the top 3 classes are: ivy.array([0.64773697, 0.29496649, 0.04526037], dev=gpu:0)
+Categories of the top 3 classes are: ['tiger cat', 'tabby', 'Egyptian cat']
+We can check to confirm that the model gets the same results as the torchvision implementation.
+
+
+
+
+
+
+[ ]:
+from torchvision.models import alexnet as torch_alexnet
+from torchvision.models import AlexNet_Weights
+
+torch_alexnet = torch_alexnet(weights=AlexNet_Weights.IMAGENET1K_V1, dropout=0)
+
+
+
+
+
+
+[ ]:
+torch_output = torch.softmax(torch_alexnet(torch_img), dim=1)
+torch_classes = torch.argsort(torch_output[0], descending=True)[:3]
+torch_logits = torch.take(torch_output[0], torch_classes)
+
+print("Indices of the top 3 classes are:", torch_classes)
+print("Logits of the top 3 classes are:", torch_logits)
+print("Categories of the top 3 classes are:", [categories[i] for i in torch_classes])
+
+
+
+
+
+
+
+Indices of the top 3 classes are: tensor([282, 281, 285], device='cuda:0')
+Logits of the top 3 classes are: tensor([0.6477, 0.2950, 0.0453], device='cuda:0', grad_fn=<TakeBackward0>)
+Categories of the top 3 classes are: ['tiger cat', 'tabby', 'Egyptian cat']
+Great! We can see that the classes and corresponding logits are the same in the Ivy native model and the torchvision model.
+TensorFlow inference#
+With an Ivy native model, it is simple to change the backend, which lets us use the model in a different framework. First we’ll try TensorFlow.
+
+
+
+
+
+
+[ ]:
+import tensorflow as tf
+
+ivy.set_backend("tensorflow")
+ivy_alexnet = alexnet()
+Once the backend is set to TensorFlow, we can use TensorFlow to do our logit post-processing.
+
+
+
+
+
+
+[ ]:
+st = time.perf_counter()
+raw_logits = ivy_alexnet(ivy.asarray(img)) # pass the image to the model
+latency = time.perf_counter() - st
+output = tf.nn.softmax(raw_logits)
+classes = tf.argsort(output[0], axis=-1, direction="DESCENDING")[:3] # get the top 3 classes
+logits = tf.gather(output[0], classes) # get the logits
+
+print("Latency:", latency)
+print("Indices of the top 3 classes are:", classes)
+print("Logits of the top 3 classes are:", logits)
+print("Categories of the top 3 classes are:", [categories[i] for i in classes.numpy().tolist()])
+
+
+
+
+
+
+
+Latency: 10.652289830999962
+Indices of the top 3 classes are: tf.Tensor([282 281 285], shape=(3,), dtype=int32)
+Logits of the top 3 classes are: tf.Tensor([0.6477362 0.29496726 0.04526032], shape=(3,), dtype=float32)
+Categories of the top 3 classes are: ['tiger cat', 'tabby', 'Egyptian cat']
+As expected, the results are identical to the Torch backend! If you had another model or postprocessing routine written in TensorFlow, Ivy makes it simple to feed one into the other, without having to (carefully) rewrite them all to one backend. It also means you can easily try out different backends to see which one is the quickest for your particular model and hardware.
+Again, we can call ivy.trace_graph to speed up inference.
+
+
+
+
+
+
+[ ]:
+ivy_alexnet = ivy.trace_graph(ivy_alexnet, args=(ivy.asarray(img),))
+Repeating the previous inference, we see that the traced model gets the same results as before, and is faster.
+
+
+
+
+
+
+[ ]:
+st = time.perf_counter()
+raw_logits = ivy_alexnet(ivy.asarray(img)) # pass the image to the model
+latency = time.perf_counter() - st
+output = tf.nn.softmax(raw_logits)
+classes = tf.argsort(output[0], axis=-1, direction="DESCENDING")[:3] # get the top 3 classes
+logits = tf.gather(output[0], classes) # get the logits
+
+print("Latency:", latency)
+print("Indices of the top 3 classes are:", classes)
+print("Logits of the top 3 classes are:", logits)
+print("Categories of the top 3 classes are:", [categories[i] for i in classes.numpy().tolist()])
+
+
+
+
+
+
+
+Latency: 0.026875037000081647
+Indices of the top 3 classes are: tf.Tensor([282 281 285], shape=(3,), dtype=int32)
+Logits of the top 3 classes are: tf.Tensor([0.6477362 0.29496726 0.04526032], shape=(3,), dtype=float32)
+Categories of the top 3 classes are: ['tiger cat', 'tabby', 'Egyptian cat']
+JAX inference#
+
+
+
+
+
+
+[ ]:
+# Overrides Jax's default behavior of preallocating 75% of GPU memory
+# Temporary fix until this is handled by ivy's graph tracer
+import os
+os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"
+
+import jax
+
+ivy.set_backend("jax")
+ivy_alexnet = alexnet()
+ivy_alexnet = ivy.trace_graph(ivy_alexnet, args=(ivy.asarray(img),))
+ivy_alexnet = jax.jit(ivy_alexnet)
+
+img_jax = jax.device_put(jax.numpy.asarray(img), device=jax.devices()[0])
+
+
+
+
+
+
+[ ]:
+# warm-up
+for _ in range(5):
+ _ = ivy_alexnet(img_jax)
+
+
+
+
+
+
+[ ]:
+st = time.perf_counter()
+raw_logits = ivy_alexnet(img_jax).data # pass the image to the model
+latency = time.perf_counter() - st
+output = jax.nn.softmax(raw_logits) # pass the image to the model
+classes = jax.numpy.argsort(-output[0])[:3] # get the top 3 classes
+logits = ivy.gather(output[0], classes) # get the logits
+
+print("Latency:", latency)
+print("Indices of the top 3 classes are:", classes)
+print("Logits of the top 3 classes are:", logits)
+print("Categories of the top 3 classes are:", [categories[i] for i in np.array(classes).tolist()])
+
+
+
+
+
+
+
+Latency: 0.0022192720000475674
+Indices of the top 3 classes are: [282 281 285]
+Logits of the top 3 classes are: ivy.array([0.64773613, 0.29496723, 0.04526032], dev=gpu:0)
+Categories of the top 3 classes are: ['tiger cat', 'tabby', 'Egyptian cat']
+We get the exact same results as before. Note again that we were able to use JAX functions in calculating the top three classes.
+Let’s try another image
+
+
+
+
+
+
+[ ]:
+!wget https://github.com/raw/unifyai/models/master/images/dog.jpg
+filename = "dog.jpg"
+# Preprocess torch image
+from torchvision import transforms
+from PIL import Image
+
+preprocess = transforms.Compose([
+ transforms.Resize(256),
+ transforms.CenterCrop(224),
+ transforms.ToTensor(),
+ transforms.Normalize(
+ mean=[0.485, 0.456, 0.406],
+ std=[0.229, 0.224, 0.225]
+)])
+torch_img = Image.open(filename)
+torch_img = preprocess(torch_img)
+torch_img = torch.unsqueeze(torch_img, 0)
+
+img = torch_img.numpy()
+img_jax = jax.device_put(jax.numpy.asarray(img), device=jax.devices()[0])
+
+
+
+
+
+
+[ ]:
+from IPython.display import Image, display
+display(Image(filename))
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+[ ]:
+st = time.perf_counter()
+raw_logits = ivy_alexnet(img_jax).data # pass the image to the model
+latency = time.perf_counter() - st
+output = jax.nn.softmax(raw_logits) # pass the image to the model
+classes = jax.numpy.argsort(-output[0])[:3] # get the top 3 classes
+logits = ivy.gather(output[0], classes) # get the logits
+
+print("Latency:", latency)
+print("Indices of the top 3 classes are:", classes)
+print("Logits of the top 3 classes are:", logits)
+print("Categories of the top 3 classes are:", [categories[i] for i in np.array(classes).tolist()])
+
+
+
+
+
+
+
+Latency: 0.006431100999861883
+Indices of the top 3 classes are: [258 104 259]
+Logits of the top 3 classes are: ivy.array([0.72447652, 0.13937832, 0.05874982], dev=gpu:0)
+Categories of the top 3 classes are: ['Samoyed', 'wallaby', 'Pomeranian']
+Note that the incorrect prediction of “wallaby” is down to the AlexNet model itself, as the torchvision version returns the same logits!
+
+
+
+
+
+
+[ ]:
+st = time.perf_counter()
+raw_logits = torch_alexnet(torch_img)
+latency = time.perf_counter() - st
+torch_output = torch.softmax(raw_logits, dim=1)
+torch_classes = torch.argsort(torch_output[0], descending=True)[:3]
+torch_logits = torch.take(torch_output[0], torch_classes)
+
+print("Latency:", latency)
+print("Indices of the top 3 classes are:", torch_classes)
+print("Logits of the top 3 classes are:", torch_logits)
+print("Categories of the top 3 classes are:", [categories[i] for i in torch_classes])
+
+
+
+
+
+
+
+Latency: 0.004749261999904775
+Indices of the top 3 classes are: tensor([258, 104, 259], device='cuda:0')
+Logits of the top 3 classes are: tensor([0.7245, 0.1394, 0.0587], device='cuda:0', grad_fn=<TakeBackward0>)
+Categories of the top 3 classes are: ['Samoyed', 'wallaby', 'Pomeranian']
+Appendix (Ivy code for AlexNet implementation)#
+As promised, here is the Ivy native source code for the AlexNet model used in this demo.
+
+
+
+
+
+
+[ ]:
+class AlexNet(ivy.Module):
+ """An Ivy native implementation of AlexNet"""
+
+ def __init__(self, num_classes=1000, dropout=0, v=None):
+ self.num_classes = num_classes
+ self.dropout = dropout
+ super(AlexNet, self).__init__(v=v)
+
+ def _build(self, *args, **kwargs):
+ self.features = ivy.Sequential(
+ ivy.Conv2D(3, 64, [11, 11], [4, 4], 2, data_format="NCHW"),
+ ivy.ReLU(),
+ ivy.MaxPool2D(3, 2, 0, data_format="NCHW"),
+ ivy.Conv2D(64, 192, [5, 5], [1, 1], 2, data_format="NCHW"),
+ ivy.ReLU(),
+ ivy.MaxPool2D(3, 2, 0, data_format="NCHW"),
+ ivy.Conv2D(192, 384, [3, 3], 1, 1, data_format="NCHW"),
+ ivy.ReLU(),
+ ivy.Conv2D(384, 256, [3, 3], 1, 1, data_format="NCHW"),
+ ivy.ReLU(),
+ ivy.Conv2D(256, 256, [3, 3], 1, 1, data_format="NCHW"),
+ ivy.ReLU(),
+ ivy.MaxPool2D(3, 2, 0, data_format="NCHW"),
+ )
+ self.avgpool = ivy.AdaptiveAvgPool2d((6, 6))
+ self.classifier = ivy.Sequential(
+ ivy.Dropout(prob=self.dropout),
+ ivy.Linear(256 * 6 * 6, 4096),
+ ivy.ReLU(),
+ ivy.Dropout(prob=self.dropout),
+ ivy.Linear(4096, 4096),
+ ivy.ReLU(),
+ ivy.Linear(4096, self.num_classes),
+ )
+
+ def _forward(self, x):
+ x = self.features(x)
+ x = self.avgpool(x)
+ x = ivy.reshape(x, (x.shape[0], -1))
+ x = self.classifier(x)
+ return x
+
+
+
+
+
+
+
+
+
+
+