For a better experience, open in Colab: 
For a better experience, open in Colab:
This tutorial walks you through the core ideas behind NOOA (NVIDIA Object-Oriented Agents). We’ll build a
BaristaAgentthat recommends drinks to sleepy customers. Along the way we’ll see that a NOOA agent is just a Python object — you add tools by adding methods, spin up new agents by instantiating the class, and strongly type its outputs like any regular Python function.
We’re not going to sell you on a whole new paradigm — you won’t walk away having learned some exotic new way to build software. What you will walk away with is plain Python, plus a small sprinkle of magic.
Install NOOA:
pip install nooa
NOOA is compatible with any LiteLLM-supported model — hosted or local. For hosted providers, you’ll need an API key. Local providers (Ollama, vLLM, any OpenAI-compatible endpoint) need no API key, just an api_base.
Pick a provider below. Replace "your-api-key" with a real key for hosted providers; local providers don’t need one.
from nooa.unifiedllm.registry import get_llm_client
# model = get_llm_client("claude-haiku-4-5", api_key="your-api-key") # Anthropic
# model = get_llm_client("ollama_chat/qwen3:1.7b", api_base="http://localhost:11434") # Ollama (local, no key)
# model = get_llm_client("hosted_vllm/Qwen/Qwen3-1.7B", api_base="http://localhost:8000/v1") # vLLM (local, no key)
model = get_llm_client("gpt-5.5", api_key="your-api-key") # OpenAI
Let’s define our first agent. In NOOA, you build an agent by subclassing Agent — the only required argument is the LLM that will power it. Here is a complete, working agent. It’s small enough to read every line.
from nooa import Agent
class BaristaAgent(Agent, llm=model):
"""You are a friendly barista at a small neighborhood cafe.""" # this becomes the agent's system prompt
async def recommend_drink(self, customer_request: str) -> str:
"""Recommend a single drink to the customer based on what they just told you.
Be warm and concise — one string sentence is plenty."""
...
Instantiate and call it. Notice that we await the method — generation methods are async.
barista = BaristaAgent()
result = await barista.recommend_drink("Hi, I feel so sleepy, I need caffeine ☕️.")
print(result)
You sound like you need a classic double espresso—bright, bold, and ready to wake you right up.
That’s it — a friendly recommendation, ready to serve. So what just happened?
Behind the scenes, NOOA does the work that keeps the interface this Pythonic. A few things worth noticing:
...) is how you tell NOOA “this method is agentic — hand it off to the LLM instead of running it as regular Python”Concretely, NOOA quietly assembles a prompt that looks roughly like:
<system prompt>
You are a friendly barista at a small neighborhood cafe.
<available methods>
recommend_drink(customer_request: str) -> str
Recommend a single drink to the customer based on what they just told you. Be warm and concise — one sentence is plenty.
</available methods>
Plus a few other pieces we’ll unpack later. Let’s call the agent one more time, just for fun:
result = await barista.recommend_drink("I would love something that works with a cornetto, but as you know, no capuccinos after 11am.")
print(result)
A creamy cappuccino would be lovely with a cornetto—soft foam, gentle espresso, and just the right breakfast-café feel.
📝 Takeaway: the agent is an object.
The barista just offered a strong caffeine kick — but it’s 9pm and a double espresso is definitely not the move. How do we teach the agent to respect a “no caffeine after 4pm” policy?
In most agent frameworks, you’d register a tool, describe it in a JSON schema, and wire it into the runtime. In NOOA, you just add a method to the class. Anything the agent can see on itself, it can call as a tool.
Let’s add our first tool to BaristaAgent.
from datetime import datetime
from nooa import Agent
class BaristaAgent(Agent, llm=model):
"""You are a friendly barista at a small neighborhood cafe."""
def is_only_decaf_hour(self) -> bool:
"""Return True if we should only be serving decaf right now. After 2pm we go decaf-only so our customers can still sleep tonight."""
return datetime.now().hour >= 14
async def recommend_drink(self, customer_request: str) -> str:
"""Recommend a single drink to the customer based on what they just told you.
Be warm and concise — one string sentence is plenty."""
...
barista = BaristaAgent()
await barista.recommend_drink("Hi, I feel so sleepy.")
'Since it's decaf-only right now, I'd make you a cozy decaf latte to perk up your mood without keeping you up later.'
Two things to sit with:
is_only_decaf_hour anywhere. No @tool, no JSON schema, no tools=[...] list. The framework rendered doc(self) into the system prompt (that <self> block you saw earlier), and the LLM discovered the method the same way you’d discover it while reading someone else’s code.recommend_drink’s docstring mentions is_only_decaf_hour. The LLM spotted a method that looked relevant, called it, and folded the result into its recommendation. That’s the whole “tools are just methods” idea in a single page.📝 Takeaway: in NOOA, ordinary Python methods and agentic methods live side by side on the same class. The agent freely calls the deterministic ones as tools — no registration, no schema, no glue code.
What if our cafe only serves a fixed menu, and we want the agent to only recommend drinks we actually offer?
Most agent frameworks deal in a single data type: text. Text gets passed to tools, text gets exchanged between agents, text comes back as output — and then you painstakingly parse it into JSON, hoping the model got the shape right (which, even with the best models, it doesn’t always). NOOA takes a different approach: because the agent lives inside a Python program, we can strongly type everything.
Let’s put a real type on the return value of recommend_drink:
from enum import Enum
from nooa import Agent
# Define a type for the return value
class Drink(Enum):
ESPRESSO = "espresso"
CAPPUCCINO = "cappuccino"
FLAT_WHITE = "flat white"
class BaristaAgent(Agent, llm=model):
"""You are a friendly barista at a small neighborhood cafe."""
def is_only_decaf_hour(self) -> bool:
"""Return True if we should only be serving decaf right now. After 4pm we go decaf-only so our customers can still sleep tonight."""
return datetime.now().hour >= 16
async def recommend_drink(self, customer_request: str) -> tuple[str, Drink]:
"""Pick the single best drink for the customer from the menu, based on what they told you."""
...
barista = BaristaAgent()
reason, drink = await barista.recommend_drink("Hi, I feel so sleepy.")
print("Barista: ", reason)
print("Recommended drink: ", drink)
print(type(drink))
Barista: It's decaf-only right now, but I'd make you a cozy decaf cappuccino.
Recommended drink: Drink.CAPPUCCINO
<enum 'Drink'>
This isn’t just prompt engineering — NOOA enforces the return type at runtime. If the LLM returns something that isn’t a valid Drink, the framework retries until it produces one. You get real Python objects back, not strings you have to reparse.
# To see what "validation" actually means, try constructing a Drink from
# something that isn't on the menu. Python's Enum machinery rejects it:
try:
Drink("matcha")
except ValueError as e:
print(f"ValueError: {e}")
# When the LLM returns a value that doesn't match the declared type, NOOA
# catches this same error, feeds it back into the next turn as a hint, and
# asks the model to try again. You never see the invalid value in your code.
ValueError: 'matcha' is not a valid Drink
📝 Takeaway: strong typing all the way through.
Our cafe has a finite stash of coffee beans, and every drink burns through a few. Since our agent is just a Python object, giving it state is as easy as adding a field in __init__:
from nooa import Agent
from typing import Union
class Drink(Enum):
ESPRESSO = "espresso"
CAPPUCCINO = "cappuccino"
FLAT_WHITE = "flat white"
TEA = "tea"
class BaristaAgent(Agent, llm=model):
"""You are a friendly barista at a small neighborhood cafe."""
def __init__(self, coffee_beans: int) -> None:
super().__init__()
self.coffee_beans = coffee_beans
def is_only_decaf_hour(self) -> bool:
"""Return True if we should only be serving decaf right now. After 6pm we go decaf-only so our customers can still sleep tonight."""
return datetime.now().hour >= 18
def serve(self, drink: Drink) -> Drink:
"""Serve a drink. Deducts 100 beans unless it's tea."""
if drink != Drink.TEA:
self.coffee_beans -= 100
return drink
async def recommend_drink(self, customer_request: str) -> tuple[str, Union[Drink, None]]:
"""Recommend a single drink to the customer based on what they just told you.
Be warm and concise — one sentence is plenty.
We currently have {self.coffee_beans} beans left. If we ran out, recommend a tea.
Call self.serve(drink) with your chosen drink before returning."""
...
barista = BaristaAgent(coffee_beans=200)
for _ in range(3):
barista_answer, drink = await barista.recommend_drink("Something to keep me going, please.")
print(f"Answer: {barista_answer}\nServed: {drink}\nBeans left: {barista.coffee_beans}")
Answer: Absolutely — an espresso should give you a nice little boost.
Served: Drink.ESPRESSO
Beans left: 100
Answer: Absolutely — an espresso should give you a nice little boost.
Served: Drink.ESPRESSO
Beans left: 0
Answer: I'm out of coffee beans just now, but a bright cup of tea will still keep you nicely refreshed.
Served: Drink.TEA
Beans left: 0
📝 Takeaway: the agent “sees” everything on itself — including itself.
NOOA doesn’t hide the prompt from you. print_prompt renders exactly what would be sent to the LLM for a given method call — the system prompt, the agent introspection block, and the task. Take a look:
import nooa
await nooa.print_prompt(barista.recommend_drink, customer_request="stressed and running late")
Look through the output. There’s no hidden state — this is exactly what the LLM sees. A few blocks worth naming:
<system_prompt> — opens with your class docstring. This is the persona the model wears for every method on the class.<strategy_prompt> — a compact rulebook for how to act each turn: what tools exist (execute_python, return_result), when to use which, and how to finish a run. This is what teaches the LLM to “inhabit” the framework, and it’s the same for every agent.<execution_context> — the imports, types, and helpers that will be in scope when the LLM writes code. Anything you import at module level shows up here.<self> — auto-generated documentation of the agent’s public methods and fields, rendered from doc(type(self)). This is how the LLM discovers what the agent can do — no separate tool registry needed.That’s the whole prompt. No hidden system messages, no template files, no per-tool JSON schemas glued on the side. The reason the “rulebook” stays short is that the framework is plain Python, and the model already knows Python.
Coming later:
print_promptonly shows the outgoing prompt. A later notebook introduces the live trace viewer for watching an entire run unfold — the LLM response, generated code, helper calls, retries, validation, and final return value.
Five things the barista taught us:
... marks a generation method. No decorator, no separate registry. If the body is ..., the LLM implements it.{self.attr} in docstrings is live. Change the attribute, and the next call sees the new value.self is a tool. No @tool decorator, no JSON schema, no registration step.<mask> tokens given context, effectively learning a denoising task.
In this post, we’ll explore how DLMs work “under the hood” using one of the most elegant interpretability tools out there: LogitLens. The final result will be a plot like this one, that we obtained applying LogitLens to Llama in a previous blog post.
As usual, here is the code to reproduce all the plots, just
LogitLens is a clever interpretability method originally designed for autoregressive (AR) LLMs. Instead of applying the final language modeling head only to the last hidden layer (as is standard), LogitLens applies it at every intermediate layer. This lets us “see” how predictions evolve across layers—like watching a thought form in real time.
In a previous blog, I showed how to implement LogitLens for LLMs from scratch. This time, we’ll use it to analyze DLMs—specifically LLaDA and Dream—to observe how they construct predictions layer by layer. I should mention here that this is not the most advanced interpretability method (e.g. see this nice blog) but it still gives us cool insights into the model’s inner workings!
Let’s start by loading the LLaDA-8B model and giving it a classic prompt:
import torch
from transformers import AutoTokenizer, AutoModel
device = 'cuda'
model_name = 'GSAI-ML/LLaDA-8B-Base'
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.bfloat16).to(device).eval()
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
Now we create a simple prompt:
example = "The quick brown fox jumps"
prompt = tokenizer(example, return_tensors="pt").to(device).input_ids
Unlike GPT-style models that generate one token at a time, diffusion models generate N tokens simultaneously, where N is the number of masked tokens in the input sequence. This means they need to know the amount of masked tokens, that is, their target length, before they begin. We can think of it like filling out a crossword puzzle—you need to know exactly how many squares you’re working with before you can start placing letters. This might sound like a design problem for open ended generation (it kind of is) but we can circumvent with various strategies.
To prepare our input, we create a sequence with the original prompt + 8 mask tokens <mask>. The <mask> tokens are the squares of the puzzle the model has to fill in.
mask_id=126336 # specific to llada model
gen_length=8
x = torch.full((1, prompt.shape[1] + gen_length), mask_id, dtype=torch.long).to(model.device)
x[:, :prompt.shape[1]] = prompt.clone()
Just to make sure we see what we are doing, let’s visualize what is happening and print the input before feeding it to the model:
tokenizer.convert_ids_to_tokens(x[0], skip_special_tokens=False)
['The', 'quick', 'brown', 'fox', 'jumps', '<|mdmmask|>', '<|mdmmask|>', '<|mdmmask|>', '<|mdmmask|>', '<|mdmmask|>', '<|mdmmask|>', '<|mdmmask|>', '<|mdmmask|>']
It makes sense, we have our prompt + the mask tokens the model will try to fill in.
Let’s forward the input through the model and grab the predicted tokens:
# we need all the intermediate hidden states
with torch.no_grad():
outputs = model(x, output_hidden_states=True)
# get the predicted tokens and decode them
predicted_token_ids = outputs["logits"].argmax(-1)
tokenizer.convert_ids_to_tokens(predicted_token_ids[0], skip_special_tokens=False)
['_', 'quick', 'brown', 'fox', 'jumps', 'over', 'the', 'lazy', 'dog', '.', '_', '_', 'jumps']
Well, after just one decoding step we were not expecting any special results. Usually, we perform at least N/2 denoising steps, where N is the generated sequence length. However, we see the model has already guessed (correclty) almost all the missing words. On the other hand, the first token The has been replaced by a _, which we don’t mind as being part of the prompt we would just ignore that during decoding. If you want to get a better prediction, you should perform additional diffusion steps, by remasking some of the input tokens. This is out of the scope for this blog, as we just want to explore the inner representations of the model.
We just got the predictions fot the final output of the model, that is, what comes out of the last layer. What if we apply the model head (our “Lens”) to intermediate representations though? Intermediate representations are stored in the outputs.hidden_states. We’ll have one for each layer of the model and each token of the sequence.
logitlens = []
for hidden_state in hidden_states:
hidden_state = model.model.transformer.ln_f(hidden_state)
logits = model.model.transformer.ff_out(hidden_state)
# get predictions for this layer
predicted_token_ids = logits.argmax(-1)
predicted_tokens = tokenizer.convert_ids_to_tokens(predicted_token_ids[0], skip_special_tokens=False)
# store in our logitlens list
logitlens.append(predicted_tokens)
Now the list contains decoded sequences for each layer. For instance, if we want to check what the model was “thinking” at layer 5, we can just do:
print(logitlens[5])
['paralleled', 'awaited', 'funnels', 'gfx', 'tesy', 'awaited', 'awaited', 'blockList', 'blockList', 'blockList', 'blockList', 'blockList', 'blockList']
The raw output appears chaotic and difficult to interpret. To better understand the model’s internal reasoning process, let’s visualize the complete sequence predictions across all transformer layers using a heatmap (full implementation available in the accompanying notebook).
The bottom row shows the original input tokens, while each ascending row represents successive transformer layers from early to final. The color intensity indicates entropy ( ~ prediction confidence) at each layer, with darker cells representing higher entropy (model uncertainty) and brighter cells showing lower entropy (confident predictions). – that is why brighter cells are only in the final layers.
The most evident pattern is that coherent, confident predictions only emerge in the final layers, typically around layer 25 and beyond. Throughout the majority of the network depth, the model remains highly uncertain about its predictions, as evidenced by the prevalence of darker cells (and wrong tokens) in earlier layers. Additionally, correct tokens only emerge at the very final layer. This behavior fundamentally differs from autoregressive models like LLaMA or Gemma, where confident predictions typically crystallize much earlier in the network.
This behavior might be caused by a lot of factors. Keeping in mind that logit lens can be deceiving as it depends intermediate layer’s basis vector alignment with the output space, we can conjecture that the model is performing a lot of computation that is not interpretable in intermediate layers, and creating semantic representation only in the final ones. This might be a consequence of the bidirectional attention (implies more “mixing” of the tokens hence more overlapped signals?).
Alos, we should not forget taht we’re observing only the first denoising step of what is inherently a multi-step iterative process. The model is designed to gradually refine predictions across multiple iterations, so initial uncertainty and errors are not only expected but necessary for the denoising mechanism to function properly. More precisely, the more mask tokens, the harder the task hence the uncertainty for the model.
Finally, let’s repeat the process with Dream (check the Colab for the full code). Here, I will plot fewer layers to make them more clear:
It looks like Dream gets more confident at early layers, but still, correct predictions only emerge very late. In this plot, it’s also worth to point out an interesting behavior that we didn’t observe in the previous example: DLMs ability to “peek into the future”. If we look carefully, Dream predicts tokens related to words that come later in the sequence, which would be impossible in autoregressive models, where each token only attends to previous ones. An example is the token “Guinea” emerging before “pig” (I’m assuming the sentence “this sentence is about a Guinea pig” is not present in the training corpus here). With bi-directional attention, all tokens attend to each other from the start, enabling DLMs to make globally-informed guesses during decoding!
That’s it! If you have feedback, want to chat about this topic, or find bugs in the code, feel free to reach out on Twitter or any social platform. Thanks for reading!
P.S. I am maintaining a GitHub repo: Awesome Diffusion Language Models. You might find it interesting if you liked this post!
]]>Disclaimer: unlike other posts in this blog that actually served some purpose, this is just a random idea I had and am implementing for fun. So if your question is “why should I want to visualize the vocabulary of an LLM?”, I don’t have an answer 😄
I recently realized I have never visualized the embedding space of an LLM. And while it’s perfectly reasonable to go through life without ever seeing the inside of an LLM’s vocabulary, I personally find it satisfying to make abstract things more concrete — especially when they’re hiding in 4096-dimensional space.
LLMs are trained on vast corpora of text. To process text, they first tokenize it — splitting it into units ( =tokens) — and then represent each token as a dense vector in a high-dimensional space. These vectors are parameters of the model, learned during training.
The size of the vocabulary (i.e., how many distinct tokens the model knows) is a hyperparameter. For example, LLaMA 2: ~32,000 tokens, LLaMA 3: ~128,000 tokens, LLaMA 4: up to 200,000 tokens and so on. The token embedding dimension, that is the number of elements of each vector representing a token, (e.g., 4096) is another hyperparameter.
So if we take a LLaMA 3 model with an embedding size of 4096 and a vocabulary of 128k tokens, we’re dealing with 128k points in a 4096-dimensional space — not exactly human-interpretable. But what if we reduce the dimensionality? We’ll lose information, sure, but it might still give us interesting insights.
As usual, you can also
We’ll use the Hugging Face Transformers library to load a LLaMA 2 model, but you can replace it with any other model !
from transformers import AutoTokenizer, AutoModel
import torch
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, device_map="auto", torch_dtype=torch.bfloat16).eval()
To avoid plotting all 128k tokens, we’ll just extract those that actually appear in some real text. (If you want to visualize the entire vocabulary, I’ll leave the code at the end of the post, but consider that might generate a rather large and unexplorable plot.)
from datasets import load_dataset
# Load part of the Wikitext-2 dataset (raw text)
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:20%]")
corpus = " ".join(dataset["text"])
We tokenize the corpus and count how often each token appears. This helps us later if we want to color or size points by frequency.
from collections import Counter
# Tokenize the corpus
input_ids = tokenizer(corpus, return_tensors="pt")["input_ids"].flatten()
unique_ids = sorted(set(input_ids.tolist()))
freqs = Counter(input_ids.tolist())
Let’s extract the embeddings from the model and use PCA to reduce their dimension. This means that each token embedding will go from 4096 -> 3 dimensions.
from sklearn.decomposition import PCA
# Extract embeddings for tokens in corpus
embedding_weights = model.embed_tokens(torch.tensor(unique_ids).unsqueeze(0).to(model.device)).squeeze(0)
# Reduce to 3D using PCA
reduce_dims = PCA(n_components=3)
embeddings_3d = reduce_dims.fit_transform(embedding_weights.cpu().detach().float())
We tokenize the corpus and count how often each token appears. This helps us later if we want to color or size points by frequency.
# Decode each token to its string representation
tokens = [tokenizer.decode([i]) for i in unique_ids]
sizes = [freqs[i] for i in unique_ids]
Now let’s put it all together in an interactive 3D plot. I’m clipping the tokens that appear more than 100 times because they are outliers that would cause all the other points to have the same color.
import plotly.express as px
fig = px.scatter_3d(
x=embeddings_3d[:, 0],
y=embeddings_3d[:, 1],
z=embeddings_3d[:, 2],
hover_name=tokens,
color={k: v if v < 100 else 100 for k, v in freqs.items()}, # clip outlier freqs
title="Token Embeddings (Filtered by Corpus)",
)
fig.update_traces(marker=dict(size=3, showscale=False))
fig.write_html("my_token_embeddings.html") # save it to visualize locally
fig.show()
And this is the results:
It’s really interesting to surf the plot and explore where learned to position different tokens in the embedding space! For example, on the top right corner there are a lot of words at the top (high z coordinate) seem to be names of cities or states: Switzerland, Connecticut, irmingham and so on …
# download the corpus from https://dmf.unicatt.it/~della/pythoncourse18/commedia.txt
import requests
url = "https://dmf.unicatt.it/~della/pythoncourse18/commedia.txt"
response = requests.get(url)
with open("commedia.txt", "w") as f:
f.write(response.text)
dataset = load_dataset("text", data_files={"train": "commedia.txt"})
corpus = " ".join(dataset["train"]["text"])
Try plotting the whole vocab, this will generate a large html plot.
Uncomment at your own risk — slow and memory intensive!
full_ids = list(range(tokenizer.vocab_size))
full_embedding_weights = model.embed_tokens(torch.tensor(full_ids).unsqueeze(0).to(model.device)).squeeze(0)
full_embeddings_3d = PCA(n_components=3).fit_transform(full_embedding_weights.cpu().detach().float())
full_tokens = [tokenizer.decode([i]) for i in full_ids]
fig_full = px.scatter_3d(
x=full_embeddings_3d[:, 0],
y=full_embeddings_3d[:, 1],
z=full_embeddings_3d[:, 2],
hover_name=full_tokens,
title="Token Embeddings (Full Vocabulary)"
)
fig_full.show()
This was just a fun exercise to make LLM internals a bit more tangible. Sure, PCA compresses a lot of high-dimensional nuance into 3D, but it’s still fascinating to see structure emerge. Try zooming into different clusters, or color by other features (length, part-of-speech, etc.).
]]>Phi-1.5. LogitLens is a straightforward yet effective interpretability method.
The core idea behind it is to apply the language model’s output layer (also known as the “unembedding matrix” or “language modeling head”) to the hidden states at each layer of the transformer. This allows us to see how the model’s internal representations change as the input progresses through the network. Surprisingly, the model often acquires a significant amount of semantic understanding in the earlier layers of the transformer. By inspecting the predicted tokens at each layer, we can observe how the model’s understanding of the input evolves.
Disclaimer: ✋ If you’re looking for advanced interpretability tools, there are plenty of powerful libraries out there. But here, we’re going back to basics and do this from scratch because it’s always cool to understand how things work under the hood.
You can also
We’ll use Microsoft Phi-1.5 here since it’s a small, open model. Feel free to swap in another Hugging Face model.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_id= "microsoft/phi-1.5"
# load the model
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16).eval().to(device)
tokenizer = AutoTokenizer.from_pretrained(model_id, add_bos_token=True, bos_token='<bos>', use_fast=False)
Downloading the model might take a while, so you better pick a small model :). Let’s now consider an example input sentence and tokenize it.
example = "The quick brown fox jumps over the lazy"
inputs = tokenizer(example, return_tensors="pt").to(device)
print("Input shape: ", inputs["input_ids"].shape)
Input shape: torch.Size([1, 9])
The sentence was encoded into 9 tokens. In case you want to know what the tokens looks like, you can just decode them back:
original_input_tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0], skip_special_tokens=False)
print("Input tokens: ", original_input_tokens)
Input tokens: ['<bos>', 'The', 'Ġquick', 'Ġbrown', 'Ġfox', 'Ġjumps', 'Ġover', 'Ġthe', 'Ġlazy']
As we can see, the tokenizer added the beggining of sentence <bos> token. The ugly Ġ represent spaces.
In the notebook you can find a function clean these up a bit (I find the Ġs are really annoying):
original_input_tokens = cleanup_tokens(original_input_tokens)
original_input_tokens
['<bos>', 'The', ' quick', ' brown', ' fox', ' jumps', ' over', ' the', ' lazy']
Now, let’s feed the input into the model to get the next token prediction along with all the hidden states. Fortunately, the model’s forward method provides an option to return its hidden states.
# we need all the intermediate hidden states
with torch.no_grad():
outputs = model(**inputs, output_hidden_states=True)
# # print(outputs.keys())
print("Logits shape: ", outputs["logits"].shape)
Logits shape: torch.Size([1, 9, 51200]) # (batch, sequence len, vocab size)
The logits have been already projected into the vocabulary space. Hidden states on the other hand are still “raw” token representations. We’ll have one hiddent state vector for each model layer.
hidden_states = outputs.hidden_states
print("Number of model layers", len(hidden_states))
print("Hidden states for first layer", hidden_states[0].shape)
Number of model layers: 25
Hidden states for first layer: torch.Size([1, 9, 2048])
As we see, each layer in the model produces a hidden state. Here the last dimension represents the embedding size (not the vocbulary size).
By applying the language modeling head (or unembedding matrix) to the hidden state at any layer, we can generate ‘early’ logits—predictions from intermediate representations. While the model isn’t explicitly trained to produce meaningful logits at these layers, we’ll see that it naturally starts embedding token-level information along the way.
We can apply the language modeling head like this:
logits_at_second_layer = model.lm_head(hidden_states[2])
print("Logits at second layer shape: ", logits_at_second_layer.shape)
Logits at second layer shape: torch.Size([1, 9, 51200])
We now want to access the hidden state at each layer, apply the language modeling head to get the logits, and finally decode the logits into tokens.
for i, hidden_state in enumerate(hidden_states):
# apply the language model head to the hidden states
logits = model.lm_head(hidden_state)
# decode the logits to get the predicted token ids
predicted_token_ids = logits.argmax(-1)
# convert the token ids to tokens
predicted_tokens = tokenizer.convert_ids_to_tokens(predicted_token_ids[0], skip_special_tokens=False)
predicted_tokens = cleanup_tokens(predicted_tokens)
# append the predicted tokens to the list for later
logitlens.append(predicted_tokens)
print(f"Layer {i}: {predicted_tokens}")
Layer 0: ['-', ' S', '-', '-', '-', '-', '-', ' S', '-']
Layer 1: ['ed', 'oret', 'est', 'ies', 'es', 'uit', ' the', ' same', ' double']
Layer 2: ['import', 'oret', 'est', 'ies', 'es', 'uits', ' time', ' same', ' part']
Layer 3: ['import', 'orem', 'est', 'arf', 'es', 'es', ' all', ' entire', ' man']
Layer 4: [' realise', 'orem', 'est', 'arf', 'es', 'uit', 'worked', ' entire', ' man']
Layer 5: [' realise', 'orem', 'est', 'arf', 'es', 'uit', 'worked', ' entire', ' man']
Layer 6: [' realise', 'orem', 'est', 'arf', 'es', 'uit', 'kill', 'ses', ' man']
Layer 7: ['iveness', 'orem', 'est', ' fox', 'es', 'ers', ' all', ' entire', ' brown']
Layer 8: ['iveness', 'orem', 'ness', ' fox', 'es', 'ers', 'ind', ' entire', ' poor']
Layer 9: ['iveness', 'orem', 'ness', ' fox', 'es', 'ers', ' obstacles', ' entire', ' poor']
Layer 10: ['iveness', 'orem', 'ness', ' ph', 'es', 'ers', ' obstacles', ' entire', ' poor']
Layer 11: ['iveness', 'orem', ' brown', ' fox', 'es', 'ers', ' obstacles', ' entire', ' poor']
Layer 12: ['iveness', 'oret', ' brown', ' fox', 'es', ' into', ' obstacles', ' entire', ' poor']
Layer 13: ['ality', 'oret', ' brown', ' fox', 'es', ' into', ' obstacles', ' entire', ' poor']
Layer 14: ['ality', 'ory', ' brown', ' ph', 'es', ' into', ' obstacles', ' entire', ' poor']
Layer 15: ['iveness', 'ory', ' brown', ' fox', 'es', ' into', ' obstacles', ' entire', ' poor']
Layer 16: ['import', 'ory', ' brown', ' fox', 'es', ' into', ' lazy', ' entire', ' poor']
Layer 17: ['import', 'mes', ' brown', ' fox', 'es', ' over', ' the', ' lazy', ' poor']
Layer 18: ['import', ' first', ' brown', ' fox', 'es', ' over', ' the', ' lazy', ' dog']
Layer 19: [' example', ' first', ' brown', 'Ċ', 'es', ' over', ' the', ' lazy', ' dog']
Layer 20: ['ĊĊ', ' first', ' brown', 's', 'es', ' over', ' the', ' lazy', ' dog']
Layer 21: ['ing', ' first', ' brown', ' fox', 'es', ' over', ' the', ' lazy', ' dog']
Layer 22: ['Ċ', ' first', ' brown', 'Ċ', ' jumps', ' over', ' the', ' lazy', ' dog']
Layer 23: ['Ċ', 'Ċ', ' brown', ' fox', ' J', ' over', ' the', ' lazy', ' dog']
Layer 24: ['Ċ', 'ory', ' brown', ' fox', ' jumps', ' over', ' the', ' lazy', ' dog']
As you observe, the predictions refine layer-by-layer, reflecting the model’s gradual understanding of the input. We can visualize the predictions with a heatmap:
# create a heatmap that has a row for each list in the logitlens list
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
sns.set_theme(style="white")
# just for the bkg color
intensities = np.ones((len(hidden_states), len(original_input_tokens)))
# Create heatmap
plt.figure(figsize=(20, 10))
ax = sns.heatmap(intensities[::2],
annot=cleanup_tokens(logitlens)[::2],
fmt='',
cmap='Greys',
xticklabels=original_input_tokens,
yticklabels=list(range(len(logitlens)))[::2],
cbar=False
).invert_yaxis()
Right now, our heatmap just displays the model’s top predictions (using argmax), which is fine but a bit flat. Let’s make it more interesting by incorporating model certainty into the visualization.
A good way to quantify the model’s certainity about its output is looking at the entropy of the output distribution. Let’s replace the background color of each cell with the entropy of the model when generating that token.
We’ll calculate the entropy of the output distribution, using it to color the background:
# aux function to compute the entropy from logits
def entropy_from_logits(logits):
probs = torch.nn.functional.softmax(logits, dim=-1).clamp(1e-8, 1) #avoid nans
return -torch.sum(probs * torch.log(probs), dim=-1).squeeze()
Now we can run the same code as before, this time we’ll also compute and store the entropies.
logitlens = []
entropies = []
for i, hidden_state in enumerate(hidden_states):
# apply the language model head to the hidden states
logits = model.lm_head(hidden_state)
# get the entropy of the logits
entropy = entropy_from_logits(logits).float().cpu().detach().numpy()
# decode the logits to get the predicted token ids
predicted_token_ids = logits.argmax(-1)
# convert the token ids to tokens
predicted_tokens = tokenizer.convert_ids_to_tokens(predicted_token_ids[0], skip_special_tokens=False)
predicted_tokens = cleanup_tokens(predicted_tokens)
# append the predicted tokens to the list
logitlens.append(predicted_tokens)
entropies.append(entropy)
print(f"Layer {i}: {predicted_tokens}")
Let’s now create a plot where each cell is colored based on the entropy.
# Create figure and axis
plt.figure(figsize=(20, 10))
# Create heatmap
ax = sns.heatmap(np.stack(entropies)[::2],
annot=logitlens[::2],
fmt='',
cmap='YlGnBu',
xticklabels=original_input_tokens,
yticklabels=list(range(len(logitlens)))[::2],
).invert_yaxis()
I recently came across a nice blog by Sonia Joseph, pointing out that the logit lens is a convenient way to investigate internal representations. But it can be misleading, as it depends on the layer’s basis vector alignment with the output space. Linear probes often show that decent representations can be present in much earlier layers than what the logit lens portrays.
Hope you liked this! If you have any suggestions/questios, feel free to drop me a message/email or visit my page or my twitter @devoto_alessio.
Thanks Luigi for reviewing this!
]]>However, every time I tried to use JAX, I ended up using Flax instead, which offers a kind of object oriented interface (similar to torch). While Flax is great, it introduces additional layers of abstraction that make it similar to Pytorch and therefore I ended up wondering: “why am I doing this?”. There are other great frameworks as well, with different functionalities, like equinox (maybe closer to JAX’s original nature), but they always add “another layer”.
This time, I wanted to take a taste of bare JAX and avoid external libraries or abstractions. In this implementation, I’ve built a basic Vision Transformer entirely from scratch. Although it may not be the most efficient code, my focus is to explore JAX directly and train a small model while leveraging JAX’s core features, like vmap and jit, without any external frameworks.
I will cover the following topics:
jax.vmap)jax.jit)✋ If you are not interested in model initialization, you can just skip to the core part where we implement the model and train it.
You can also
In the following, I assume you are already familiar with the Vision Transformer architecture. If you are not, you can take a look here. In short, ViTs split images into patches, and treat patches as tokens (like words in NLP models), processing them using transformer layers with bidirectional (non masked) attention. In this post, we’ll build a small ViT that can train on the Imagenette dataset, and you can even run it on your local machine.
Speaking of GPUs, JAX offers seamless handling of hardware acceleration. It automatically detects and utilizes available GPUs/TPUs without requiring explicit code changes.
import jax
print("Available devices:", jax.devices()) # JAX will take care of the device placement for you
Available devices: [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]
In this notebook, we are going to use a small ViT, with the following hyperparameters:
image_size = 64
patch_size = 4
num_patches = (image_size // patch_size) ** 2
num_layers = 4 # number of transfomer layers
hidden_dim = 192 # hidden dimension of each token
mlp_dim = 192*4 # hidden dimension in the MLP
num_classes = 10 # Imagenette number of classes
num_heads = 4 # attention heads
head_dim = hidden_dim//num_heads
JAX is a fully functional framework, which means that model parameters are treated as a distinct set of numbers, existing “outside” the model itself. This gives you a nice, low-level feel for how the model works. Instead of encapsulating parameters within an object (like in torch), you’re directly manipulating a concrete set of weights along with a function that processes them.
To initialize these weights at random, we need some random primitives (just like in torch). In JAX, every call to a random primitive requires a random key, which ensures that the randomness is both explicit and controllable. This means that instead of going
a_tensor = torch.randn(tensor_shape)
you have to explicitly allocate a key first and then use it to generate a random number, like this:
key = random.PRNGKey(42)
a_tensor = random.normal(key, tensor_shape)
This is great for ML practitioners, and you know what I’m talking about if you ever had to use torch random seeding and ended up with reproducibility issues. The main reason for JAX explicitly tracking the random keys without using a global random state is that this would compromise the execution of parallel code, that is one of the main perks of JAX. You can read more about randomness in JAX here.
Let’s see what parameters we need for our ViT. Here is the list:
CLS (classification) tokenAs mentioned before, these parameters are just numbers that we can store in a dictionary like this:
# initialize vit parameters
vit_parameters = {
'patch_embed': None,
'positional_encoding': None,
'layers': [],
'final_layer_norm': None,
'head': [],
'cls_token': None
}
We now need to initialize each set of weights separately. Again, we could use a library for this, like optax, but we want to go through the process manually to better understand what’s happening under the hood.
We initialize the class token with all zeros:
import jax.numpy as jnp # this is what we use to manipulate tensors in JAX. It is supers similar to numpy
# for the class token, we just need a single vector of the same size as a token
cls_token = jnp.zeros((1,hidden_dim))
vit_parameters['cls_token'] = cls_token
For patch embedding, positional encoding, final head, and all transformer blocks we use random values (check the colab for complete code).
Each transformer block is made up of attention, mlp and layer normalization. We define a function to initialize each of these components. I’ll show just the mlp initialization here for brevity. I’ll do it using Xavier intialization, but this is not crucial and you can just use a random normal.
For the MLP, we need weights and biases for 2 layers.
def initialize_mlp(hidden_dim, mlp_dim, key):
w1_key, w2_key = random.split(key) # get new random keys from the one provided
# Xavier uniform limit for w1 and w2
limit = jnp.sqrt(6.0 / (hidden_dim + mlp_dim))
# Xavier uniform initialization for weights
w1 = random.uniform(w1_key, (hidden_dim, mlp_dim), minval=-limit, maxval=limit)
b1 = jnp.zeros(mlp_dim)
w2 = random.uniform(w2_key, (mlp_dim, hidden_dim), minval=-limit, maxval=limit)
b2 = jnp.zeros(hidden_dim)
return w1, b1, w2, b2
We are now ready to initialize all the weights in each transformer layer! Let’s create a set of parameters for each layer and store them in our dictionary.
key = random.PRNGKey(42)
key, *layer_keys = random.split(key, num_layers+1)
for i in range(num_layers):
mlp_params = initialize_mlp(hidden_dim, mlp_dim, layer_keys[i])
attn_params = initialize_attention(hidden_dim, num_heads, layer_keys[i])
ln1_params = initialize_layer_norm(hidden_dim)
ln2_params = initialize_layer_norm(hidden_dim)
vit_parameters['layers'].append((mlp_params, attn_params, ln1_params, ln2_params))
# we also have a final layer norm outside the loop
final_layer_norm_params = initialize_layer_norm(hidden_dim)
vit_parameters['final_layer_norm'] = final_layer_norm_params
Finally, we can now write the code for the transformer encoder!
One thing we quickly notice about JAX is that everything is a function — including models. This is very different from torch, where we usually look at the model as a composition of objects (nn.Modules). So we’ll write the forward pass as just a function. The ViT function will take the ViT parameters and an image as input, that is:
prediction = vit_function(image, vit_parameters)
As we can see, the parameters are “outside” of the model. Before writing the actual code for the ViT, we come to another special feature o JAX: parallelization. Thanks to JAX native vmap function, we’ll just pretend there is no batch dimension and then use vmap to automagically handle batches. This is a great improvement as we don’t have to reason in one additional dimension and there will be no need for stuff like batch,sequence,dim = input.shape (unlike torch.) So from now on, we’ll just ignore the batch dimension.
Don’t forget that each transformer block is nothing but a function over the model parameters and an input. For the MLP, we just perform an up and down projection with a Relu activation function in the middle. Notice that we will get the input parameters from the dictionary we created earlier.
def mlp(x, mlp_params):
# unpack the parameters
w1, b1, w2, b2 = mlp_params
# out = (Relu(x*w1 + b1))*w2 + b2
up_proj = relu(jnp.matmul(x, w1) + b1)
down_proj = jnp.matmul(up_proj, w2) + b2
return down_proj
Now self attention, the only catch here is to project into multiple heads and then concatenate back
def self_attention(x, attn_params):
# unpack the parameters
q_w, k_w, v_w, q_b, k_b, v_b = attn_params
n, d_k = x.shape # n and d_k are the sequence length of the input and the hidden dimension
# project the input into the query, key and value spaces
q = jnp.matmul(x, q_w) + q_b
k = jnp.matmul(x, k_w) + k_b
v = jnp.matmul(x, v_w) + v_b
# reshape to have heads
# n, (num_heads head_dim) -> (n, num_heads, headim) -> (num_heads, n, head_dim)
q = q.reshape(n, num_heads, head_dim).swapaxes(0, 1)
k = k.reshape(n, num_heads, head_dim).swapaxes(0, 1)
v = v.reshape(n, num_heads, head_dim).swapaxes(0, 1)
# perform multi-head attention
attention_weights_heads = jnp.matmul(q, jnp.swapaxes(k, -1, -2)) / jnp.sqrt(head_dim)
attention_weights_heads = jax.nn.softmax(attention_weights_heads, axis=-1)
# output projection
output = jnp.matmul(attention_weights_heads, v)
output = output.swapaxes(0,1).reshape(n, d_k)
return output
Finally, we can assemble attention and mlps into a transformer block.
def transformer_block(inp, block_params):
# unpack the parameters
mlp_params, attn_params, ln1_params, ln2_params = block_params
# attention
x = layer_norm(inp, ln1_params)
x = self_attention(x, attn_params)
res = x + inp # skip connection
# mlp
x = layer_norm(res, ln2_params)
x = mlp(x, mlp_params)
x = x + res
return x
Before feeding the first image to the model, we need one more additional step to transform an input image into a sequence of patches. To do that, we use einops, which offers a highly expressive interface to reshape tensors. Another way would be applying convolutions but here we are just using bare JAX code so we get a sequence of tokens from an image like this:
from einops import rearrange
patches = rearrange (image, 'c (h p1) (w p2) -> (h w) (p1 p2 c)', p1=patch_size, p2=patch_size)
With this, we are ready to go! The final transformer then works by:
Let’s implement these steps:
def transformer(image, vit_parameters):
# reshape image from c,h,w -> num_patches, patch_size*patch_size
patches = rearrange (image, 'c (h p1) (w p2) -> (h w) (p1 p2 c)', p1=patch_size, p2=patch_size)
# embed the patches into tokens
patches = jnp.matmul(patches, vit_parameters['patch_embed'])
# add positional encoding
patches = patches + vit_parameters['positional_encoding']
# append class token to sequence
cls_token = vit_parameters['cls_token']
patches = jnp.concatenate([cls_token, patches], axis=0)
# forward through all transformer blocks
for layer, block_params in enumerate(vit_parameters['layers']):
patches = transformer_block(patches, block_params)
# final layer norm
patches = layer_norm(patches, vit_parameters['final_layer_norm'])
# get the class token and apply the final head
patches = patches[0, :]
logits = jnp.matmul(patches, vit_parameters['head'][0]) + vit_parameters['head'][1]
return logits
Let’s test it on a random input:
sample_image = random.normal(key, (3 ,image_size, image_size))
prediction = transformer(sample_image, vit_parameters)
print("Output shape:", prediction.shape) # should be (num_classes,)
Output shape: (10,)
As you may have noticed, the random input is just an image without a batch dimension. Let’s see how we can add a batch dimension without modifying the code.
vmapAs anticipated, before jumping into training, we’ll look at one of the coolest features JAX offers: vmap. This allows you to vectorize your functions, meaning you can apply them over batches of data without writing explicit loops. In a way, it’s like automatic batching. You write a function that works on a single example, and vmap will apply it to all examples in a batch in one go.
For example, if you have a function that processes a single image, you can turn it into a function that processes an entire batch of images with just one line:
batched_fn = jax.vmap(single_image_fn)
This can come in handy when applying the model over a batch of data. This means that we can run one pass of our transformer over a batch of images very easily. Let’s try:
bsize = 5
sample_images = random.normal(key, (bsize, 3 ,image_size, image_size))
# if we apply the transformer to a batch of images, we should get a batch of logits
# but this will raise an error
# prediction = transformer(sample_images, vit_parameters)
Let’s apply vmap. We need to map each input in the batch (first dimension is 0) to all the parameters (second dimension in None).
prediction = jax.vmap(transformer, in_axes=(0, None))(sample_images, vit_parameters)
print("Prediction shape:", prediction.shape)
Prediction shape: (5, 10)
Actually, vmap can do way more than this. I recommend this blog for an overview.
Next up is the loss function. We’ll use the Cross-Entropy Loss, which is a standard choice for classification tasks.
def cross_entropy_loss(patches, vit_parameters, ground_truth):
prediction = jax.vmap(transformer, in_axes=(0, None))(patches, vit_parameters)
logs = jax.nn.log_softmax(prediction)
l = -jnp.mean(jnp.sum(ground_truth * logs, axis=-1))
return l
l = cross_entropy_loss(sample_images, vit_parameters, jnp.zeros((bsize, 10)).at[0, 1].set(1))
print("Loss:", l)
Loss: 0.49248534
For dataset loading, I’m going to steal some code from PyTorch. PyTorch’s data utilities work really well, and since this isn’t a post about data loading, we’ll skip the hassle of reinventing the wheel here.
from torchvision.datasets import CIFAR10, Imagenette
from torchvision import transforms
from torch.utils.data import DataLoader
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
train_dataset = Imagenette(
root='/home/aledev/datasets/imagenette3',
size="160px",
split='train',
download=False,
transform=transforms.Compose([transforms.Resize((image_size,image_size)), transforms.ToTensor(), transforms.Normalize(mean, std)])
)
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True)
test_dataset = Imagenette(
root='/home/aledev/datasets/imagenette3',
size="160px",
split='val',
download=False,
transform=transforms.Compose([transforms.Resize((image_size,image_size)), transforms.ToTensor(), transforms.Normalize(mean, std)])
)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)
Let’s code a simple evaluation function that loops over the test data and computes accuracy
from tqdm import tqdm
def eval(vit_parameters):
correct = 0
for(img, target) in tqdm(test_loader, desc="Eval", unit="item"):
img = jnp.asarray(img, dtype=jnp.float32)
target = jnp.asarray(target)
logits = jax.vmap(transformer, in_axes=(0, None))(img, vit_parameters)
prediction = jnp.argmax(logits, axis=-1)
correct += jnp.sum(prediction == target).item()
acc = correct / len(test_dataset)
return acc
accuracy = eval(vit_parameters)
print("Accuracy before training", accuracy)
Eval: 100%|██████████| 16/16 [00:18<00:00, 1.13s/item]
Accuracy before training 0.06394904458598726
Before we dive into training, we meet another cool feature of JAX: jit, that is, just in time compilation. One of JAX’s biggest selling points is its ability to automatically compile and optimize your code using just-in-time (JIT) compilation. With JAX, you can wrap your functions in jax.jit() to make them faster by turning them into optimized code. It’s a one-liner and can massively speed up your training loop. In essence, jit lets you write Python code, and JAX will magically optimize it behind the scenes. It’s not even hard to use, so there’s no reason not to take advantage of it!
Here’s how you can JIT-compile your training step:
@jax.jit
def train_step(params, data):
# your training step logic here
This is where JAX differs most from PyTorch. In PyTorch, you call .backward() on your loss, and it handles everything, i.e. computes loss and gradients, that you’ll find stored in your model parameters. In JAX, you need to manually compute gradients and update parameters yourself, which gives you a more hands-on experience with the inner workings of optimization.
To perform gradient descent, we’ll compute the gradient of the loss with respect to the parameters. In JAX, you can do this using the jax.values_and_grad function:
loss, grads = value_and_grad(cross_entropy_loss, argnums=1)(input, parameters, targets)
What will the gradients look like ? The gradients are just going to be a dictionary (pytree) with the same keys as the model parameters, but instead of holding the parameters, they will hold the gradients.
from jax import value_and_grad
# fake labels and images
sample_images = random.normal(key, (bsize, 3 ,image_size, image_size))
sample_target = jnp.zeros((bsize, 10)).at[0, 1].set(1)
current_loss, grads = value_and_grad(cross_entropy_loss, argnums=1)(sample_images, vit_parameters, sample_target)
print("Current loss:", current_loss)
print("Gradients:", grads.keys())
Current loss: 0.49248534
Gradients: dict_keys(['cls_token', 'final_layer_norm', 'head', 'layers', 'patch_embed', 'positional_encoding'])
We now have a dictionary of gradients that mirrors the structure of our parameters. To update the parameters, we’ll perform simple gradient descent. The update rule is:
\[\theta_{\text{new}} = \theta_{\text{old}} - \eta \cdot \nabla_\theta L(\theta)\]Where:
JAX has some great libraries for optimization, like optax, but for simplicity, we’ll just manually update the parameters using vanilla SGD. Notice that to do this we’d have to go throught the dictionary and update all values that have the same key. Fortunately, JAX has a function that does that for us: jax.tree.map.
We just have to tell the gradient descent rule:
updated_params = jax.tree.map(lambda p, g: p - 0.01 * g, vit_parameters, grads)
Putting everything together, the training step will look like this:
from jax import jit, value_and_grad
@jit
def train_step(patches, vit_parameters, target_one_hot):
# compute gradients
current_loss, grads = value_and_grad(cross_entropy_loss, argnums=1)(patches, vit_parameters, target_one_hot)
# update parameters
updated_params = jax.tree.map(lambda p, g: p - 0.01 * g, vit_parameters, grads)
return current_loss, updated_params
Finally, let’s train the model. We don’t expect any special results because we training without any optimization and with a super small model. Also, I’ll only train for 30 epochs here, but you can let it go on for longer.
import jax
from jax import value_and_grad
from jax import jit
from tqdm import tqdm
num_epochs = 20
for epoch in range(num_epochs):
progress_bar = tqdm(enumerate(train_loader), total=len(train_loader), desc=f"Epoch {epoch+1}/{num_epochs}")
#for (data, target) in tqdm(train_loader, desc=f'Train epoch {epoch}'):
for i, (data, target) in progress_bar:
# convert to numpy
data = jnp.asarray(data)
target = jnp.asarray(target)
# reshape and get one hot fot loss
target_one_hot = jax.nn.one_hot(target, num_classes)
current_loss, vit_parameters = train_step(data, vit_parameters, target_one_hot)
progress_bar.set_postfix({'loss': current_loss})
eval_acc = eval(vit_parameters)
print(f'Epoch: {epoch}, Eval acc: {eval_acc}')
Epoch 1/20: 100%|██████████| 37/37 [00:29<00:00, 1.26it/s, loss=2.1218076]
Eval: 100%|██████████| 16/16 [00:07<00:00, 2.17item/s]
Epoch: 0, Eval acc: 0.2573248407643312
Epoch 2/20: 100%|██████████| 37/37 [00:17<00:00, 2.07it/s, loss=2.162859]
Eval: 100%|██████████| 16/16 [00:07<00:00, 2.10item/s]
Epoch: 1, Eval acc: 0.25095541401273885
...
Hope you enjoyed this, please reach me at https://alessiodevoto.github.io/ if you have any questions or find inconsistencies!
Thanks to Rahil for spotting a bug in the attention shapes!
]]>CLS token in a pretrained Vision Transformer from the timm library.
For a better experience, open in Colab:
In this short notebook, we’ll try to get some insights into pre-trained vision transformers by looking at attention patterns. More specifically, we’ll plot the attention scores between the CLS token and other tokens and check whether they have a semantic interpretation or not. This is often the case, so we expect to images like this:
# install timm
!pip install timm
We load a pre-trained DeiT (data efficient Vision Transformer) see he here.
Anyway, Timm has
plenty of
pre-trained models to choose from.
import torch
model = torch.hub.load('facebookresearch/deit:main', 'deit_tiny_patch16_224', pretrained=True)
print(model)
VisionTransformer(
(patch_embed): PatchEmbed(
(proj): Conv2d(3, 192, kernel_size=(16, 16), stride=(16, 16))
(norm): Identity()
)
(pos_drop): Dropout(p=0.0, inplace=False)
(patch_drop): Identity()
(norm_pre): Identity()
(blocks): Sequential(
(12x): Block(
(norm1): LayerNorm((192,), eps=1e-06, elementwise_affine=True)
(attn): Attention(
(qkv): Linear(in_features=192, out_features=576, bias=True)
(q_norm): Identity()
(k_norm): Identity()
(attn_drop): Dropout(p=0.0, inplace=False)
(proj): Linear(in_features=192, out_features=192, bias=True)
(proj_drop): Dropout(p=0.0, inplace=False)
)
(ls1): Identity()
(drop_path1): Identity()
(norm2): LayerNorm((192,), eps=1e-06, elementwise_affine=True)
(mlp): Mlp(
(fc1): Linear(in_features=192, out_features=768, bias=True)
(act): GELU(approximate='none')
(drop1): Dropout(p=0.0, inplace=False)
(norm): Identity()
(fc2): Linear(in_features=768, out_features=192, bias=True)
(drop2): Dropout(p=0.0, inplace=False)
)
(ls2): Identity()
(drop_path2): Identity()
))
...
The original code can be found here. We can see the attention scores are not returned (unlike the Pytorch implementation) so we have to "feature extract" them.
More specifically, we see that in the attention class the attention score are computed as
attn = self.attn_drop(attn)
and only when sdpa attention is not enabled. Before going on, let's
disable the sdpa attention in each block.
for block in model.blocks:
block.attn.fused_attn = False
Now we are ready to extract the features. We'll do it with a very cool
torch feature extraction tool called torch.fx. This allows you to
extract all intermediate activations from a model without the cumbersome
process of adding hooks or subclassing the forward, you can find more
info
here.
Let's see which features we can extract, using get_graph_node_names.
from torchvision.models.feature_extraction import get_graph_node_names
nodes, _ = get_graph_node_names(model)
print(nodes)
['x', 'patch_embed.getattr', 'patch_embed.getitem', 'patch_embed.getitem_1', 'patch_embed.getitem_2', 'patch_embed.getitem_3', 'patch_embed.eq', 'patch_embed._assert', 'patch_embed.eq_1', 'patch_embed._assert_1', 'patch_embed.proj', 'patch_embed.flatten', 'patch_embed.transpose', 'patch_embed.norm', 'pos_embed', 'cls_token', 'getattr', 'getitem', 'expand', 'cat', 'add', 'pos_drop', 'patch_drop', 'norm_pre', 'blocks.0.norm1', 'blocks.0.attn.getattr', 'blocks.0.attn.getitem', 'blocks.0.attn.getitem_1', 'blocks.0.attn.getitem_2', 'blocks.0.attn.qkv', 'blocks.0.attn.reshape', 'blocks.0.attn.permute', 'blocks.0.attn.unbind', 'blocks.0.attn.getitem_3', 'blocks.0.attn.getitem_4', 'blocks.0.attn.getitem_5', 'blocks.0.attn.q_norm', 'blocks.0.attn.k_norm', 'blocks.0.attn.mul', 'blocks.0.attn.transpose', 'blocks.0.attn.matmul', 'blocks.0.attn.softmax', 'blocks.0.attn.attn_drop', 'blocks.0.attn.matmul_1', 'blocks.0.attn.transpose_1', 'blocks.0.attn.reshape_1', 'blocks.0.attn.proj', 'blocks.0.attn.proj_drop', 'blocks.0.ls1', 'blocks.0.drop_path1', 'blocks.0.add', 'blocks.0.norm2', 'blocks.0.mlp.fc1', 'blocks.0.mlp.act', 'blocks.0.mlp.drop1', 'blocks.0.mlp.norm', 'blocks.0.mlp.fc2', 'blocks.0.mlp.drop2', 'blocks.0.ls2', 'blocks.0.drop_path2', 'blocks.0.add_1', 'blocks.1.norm1', 'blocks.1.attn.getattr', 'blocks.1.attn.getitem', 'blocks.1.attn.getitem_1', 'blocks.1.attn.getitem_2', 'blocks.1.attn.qkv', 'blocks.1.attn.reshape', 'blocks.1.attn.permute', 'blocks.1.attn.unbind', 'blocks.1.attn.getitem_3', 'blocks.1.attn.getitem_4', 'blocks.1.attn.getitem_5', 'blocks.1.attn.q_norm', 'blocks.1.attn.k_norm', 'blocks.1.attn.mul', 'blocks.1.attn.transpose', 'blocks.1.attn.matmul', 'blocks.1.attn.softmax', 'blocks.1.attn.attn_drop', 'blocks.1.attn.matmul_1', 'blocks.1.attn.transpose_1', 'blocks.1.attn.reshape_1', 'blocks.1.attn.proj', 'blocks.1.attn.proj_drop', 'blocks.1.ls1', 'blocks.1.drop_path1', 'blocks.1.add', 'blocks.1.norm2', 'blocks.1.mlp.fc1', 'blocks.1.mlp.act', 'blocks.1.mlp.drop1', 'blocks.1.mlp.norm', 'blocks.1.mlp.fc2', 'blocks.1.mlp.drop2', 'blocks.1.ls2', 'blocks.1.drop_path2', 'blocks.1.add_1', 'blocks.2.norm1', ...]
# A lot of useless stuff
# we only care for nodes that contain attn_drop
interesting_nodes = [x for x in nodes if 'attn_drop' in x]
print(interesting_nodes)
['blocks.0.attn.attn_drop', 'blocks.1.attn.attn_drop', 'blocks.2.attn.attn_drop', 'blocks.3.attn.attn_drop', 'blocks.4.attn.attn_drop', 'blocks.5.attn.attn_drop', 'blocks.6.attn.attn_drop', 'blocks.7.attn.attn_drop', 'blocks.8.attn.attn_drop', 'blocks.9.attn.attn_drop', 'blocks.10.attn.attn_drop', 'blocks.11.attn.attn_drop']
Makes sense, we have one attention for each layer.
Before going on, some standard stuff to normalize and denormalize the image for plotting.
!wget https://raw.githubusercontent.com/alessiodevoto/notebooks/refs/heads/main/data/bird.jpg
Some image processing basic stuff.
# Load and preprocess image
from PIL import Image
from torchvision import transforms
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
img = Image.open('bird.jpg')
preprocess = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std),
])
def denormalize(image):
denormalized_image = image * torch.tensor(std).view(3, 1, 1) + torch.tensor(mean).view(3, 1, 1)
return denormalized_image
img_tensor = preprocess(img).unsqueeze(0)
We finally extract the attention scores. We see we are interested in all
those nodes that contain attn_drop.
from torchvision.models.feature_extraction import create_feature_extractor
feature_extractor = create_feature_extractor(
model, return_nodes=interesting_nodes)
# `out` will be a dict of Tensors, each representing a feature map
out = feature_extractor(img_tensor)
for k, v in out.items():
print(k, v.shape)
We see the attention scores have shape
(batch, num_heads, num_patches+1, num_patches+1), where the +1 is
because we added the CLS token.
Let's iterate over the attention scores and plot them for each head.
import matplotlib.pyplot as plt
import numpy as np
num_layers = 12
num_heads = 3
# create subplots of 12 x 4
fig, axs = plt.subplots(num_layers, num_heads+1, figsize=(12, 24))
for i, (k, v) in enumerate(out.items()):
# class token attention scores
attn_scores = v.squeeze() #remove the batch dimension
# print(attn_scores.shape)
for head in range(num_heads):
attn_scores_head = attn_scores[head]
# print(attn_scores_head.shape)
cls_token_attn_scores = attn_scores_head[0,1:]
# print(cls_token_attn_scores.shape)
axs[i,head+1].imshow(cls_token_attn_scores.reshape(14,14).detach(), cmap='viridis')
axs[i,0].imshow(denormalize(img_tensor).detach().numpy().squeeze().transpose(1,2,0))
# hide ticks
for ax in axs.flat:
ax.set(xticks=[], yticks=[])
plt.tight_layout()
Hope you liked this! If you have any suggestions/questios, feel free to drop me a message/email or visit my page or my twitter @devoto_alessio.
]]>Obviously, all these methods assume you have multiple GPUs at your disposal (surprise!).
TLDR: Split your dataset across multiple GPUs, each with a full model copy. Synchronize gradients after each pass.
Data parallelism is simple to implement, and scales well with number of devices for smaller model. It is especially effective for large datasets. On the other hand, it introduces a lot of communication overhead for gradient synchronization. Additionally, because a full copy of the model is stored on each device, it also causes memory redundancy.
Pseudocode:
# On each device
for batch in dataloader:
outputs = model(batch)
loss = criterion(outputs, targets)
loss.backward()
# Synchronize gradients across devices
all_reduce(model.parameters.grad)
optimizer.step()
TLDR: Divide your model across devices, each processes the same input at different stages.
Model parallelism is perfect for handling models too large for a single device. In doing so, it also reduces the memory required for a single device. Unfortunately, it might be complex to implement efficiently, because of potential load imbalance: it usually needs pipelining to avoid GPUs from remaining idle.
Pseudocode:
# Define model portions
model_part1 = nn.Sequential(layer1, layer2).to('cuda:0')
model_part2 = nn.Sequential(layer3, layer4).to('cuda:1')
# Forward pass
def forward(x):
x = model_part1(x)
x = x.to('cuda:1')
return model_part2(x)
TLDR: Split your model into stages on different devices. Data flows through the pipeline, with multiple batches processed simultaneously.
This is the solution to the imbalancing problem for plain model parallel. A nice explanation of model parallel + pipeline parallel can be found here. Pipeline parallel balances computation and communication and makes model parallelism more efficient. As a drawback, it requires a potentially complex scheduling: you have to deal with splitting the input across GPUs and schedule the pipeline. If you do this the wrong way, you might cause “bubble” periods of idle time.
Pseudocode:
# Define stages
stage1 = nn.Sequential(layer1, layer2).to('cuda:0')
stage2 = nn.Sequential(layer3, layer4).to('cuda:1')
# Pipeline forward
def pipeline_forward(batches):
for i, batch in enumerate(batches):
x = stage1(batch)
x = x.to('cuda:1')
if i > 0:
yield stage2(prev_x)
prev_x = x
yield stage2(prev_x)
TLDR: Partition individual tensors (weights, activations) across devices. Each computes a portion of tensor operations.
This is somewhat on another level of abstraction wrt to Data and Model parallel, as tensor can represent anything in a deep learning pipeline. In other words, tensor parallel includes model and data parallel.
Pseudocode:
# Simplified tensor parallel linear layer
class TPLinear(nn.Module):
def __init__(self, in_features, out_features, n_devices):
super().__init__()
self.weight = nn.Parameter(torch.randn(out_features // n_devices, in_features))
def forward(self, x):
local_out = F.linear(x, self.weight)
return all_gather(local_out)
TLDR: Shards model parameters, gradients, and optimizer states across devices.
ZeRO includes all types of parallelism. More specifically, it impleemnts three possible options:
ZeRO offers three progressive levels of memory optimization. Each level increases memory efficiency but also introduces more communication overhead. ZeRO-3 provides the highest memory efficiency but with the most communication. More specifically:
ZeRO offers the most flexibility by combining benefits of data and model parallelism. Obviously, it introduces increased communication overhead and its complexity increases with higher ZeRO levels. Implementation of ZeRO is typically used through libraries like DeepSpeed or PyTorch’s FSDP.
For instance, latency and throughput are critical for evaluating how well a model performs in real-time applications during inference. On the other hand, floating-point operations per second (FLOPs) and parameter size are often considered during the training phase to gauge the computational and memory demands of a model.
In this short post, I will recap the most important metrics used to measure efficiency of models and highlight their differences.
The number of parameters in a model refers to the total count of learnable weights. A higher number of parameters can lead to a more powerful model capable of capturing complex patterns. However, it also demands more memory, which can make training and deployment challenging. In the past recent years, large language models (LLMs) have become extremely popular. LLMs can contain huge number of parameters, often running into billions.
Here we provide a simple table containing the number of parameters in most popular neural network layers, assuming we have an input with \(c_i\) channels, and an output with \(c_o\) channels. For convolutions, we denote the height and width of the kernel with \(k_h, k_w\) respectively.
| Layer | Number of Parameters (bias is ignored) | Explanation |
|---|---|---|
| Linear Layer | \(c_o \cdot c_i\) | simply the number of edges in a fully connected |
| Convolution | \(c_o \cdot c_i \cdot k_h \cdot k_w\) | for each input channel we have a kernel \(k_h \times k_w \times c_o\) |
| Grouped Convolution | \(\frac{c_o}{g} \cdot \frac{c_i}{g} \cdot k_h \cdot k_w \cdot g = c_o \cdot c_i \cdot k_h \cdot k_w / g\) | we group convolutions into \(g\) groups |
| Attention Block | \(3 \cdot c_i \cdot c_o + c_o \cdot c_o\) | we first project into \(QKV\) and then perform output prjection |
It is important to point out that the number of parameters does not coincide with the actual memory required to store the model, as this depends on the floating point precision. Floating point precision determines how many bits we use to store each parameter. As a consequence, the final size of the model can be computed as:
\[\text{Parameters size} = \text{number of parameters} \times \text{parameter size}\]We often try to use as few bits as possible to represent model weights, which is called quantization. Quantization reduces the precision of the model’s parameters, typically from 32-bit floating-point to 16-bit or even 8-bit integers, significantly decreasing the memory footprint and computational requirements.
As an example, the size of llama3-8b in standard fp16 would be roughly \(8000000000 \times 16 bits \sim 16 \text{GB}\). If we employ int4 quantization, the size goes down to \(8000000000 \times 4 bits \sim 4 \text{GB}\) !
The parameters size is independent from the underlying hardware, and impacts the memory during training (we have parameter efficient fine-tuning methods to tackle that) and inference (quantization helps a lot here). Additionally, as a general rule, a model with fewer parameters generally requires less compute, so parameters size often indirectly affects computational demands.
MACs (Multiply Add operations) and FLOPs (Floating Point Operations) measure the computational effort required to execute a function. One MAC is defined as:
\[a = a \times b + c\]One MAC requires performing one multiplication and one addition, i.e. two generic Floating Point Operations. Hence, we have that \(FLOPs = 2 \times MACs\).
MACs and FLOPs provide a hardware-independent way to estimate the computational cost, allowing comparisons across different models and architectures. Lowering the FLOPs while maintaining model performance is a common goal, as it can lead to faster inference times and lower energy consumption, making the model more suitable for deployment in resource-constrained environments.
However, it’s important to note that lower FLOPs do not necessarily translate to lower latency. For example, a model might have fewer FLOPs but require more memory access operations, which can be slower than the arithmetic computations themselves. Conversely, a model with higher FLOPs might be highly optimized for parallel processing, leading to lower latency on specific hardware. Thus, while FLOPs are a valuable metric for assessing computational cost, they should be considered alongside other metrics like latency to get a comprehensive view of a model’s efficiency.
In neural networks we perform a lot of matrix - matrix multiplications. For a matrix-matrix multiplication between \(A_{m \times n} \times B_{n \times k}\) we need \(nmk\) MACs and \(2nmk\) FLOPs. You can use this tool to count the FLOPs of a model in Pytorch. Let’s take a look at the FLOPs required by each of the most common neural networks layers.
| Layer | MACs | Explanation |
|---|---|---|
| Linear Layer | \(c_o \cdot c_i\) | vector matrix multiplication |
| Convolution | \(c_o \cdot o_w \cdot o_h \cdot c_i \cdot k_h \cdot k_w\) | for each output pixel, we perform the convolution \(k_h \times k_w \times c_i\) |
| Grouped Convolution | \(\frac{c_o}{g} \cdot o_w \cdot o_h \cdot c_i \cdot k_h \cdot k_w\) | we group convolutions into \(g\) groups |
| Attention Block | \(3 \cdot N \cdot {c_i}^2 + N^2 \dot c_i + N^2 \cdot {c_i}^2\) | QKV projection + attention computation + output projection |
FLOPs determine the computational intensity of model, and are the factor that most prominently affects the compute resources. The number of FLOPs that a the specific hardware can perform in a second, i.e. FLOPs per Second, is called FLOPS and measures hardware performance.
A side note on matrix multiplication. Modern GPUs are super-optimized for matrix multiplication, while other operations (e.g. layer normaliazions et similia) are less optimized. Still, this is not a problem as modern deep learnign models rely a lot on matrix multiplications. As a case study, this is how the FLOPs for a Bert model are distributed among classes of operations (credit):
Latency refers to the time it takes for a model to process an input and produce an output. This metric is particularly important in real-time applications, where quick responses are critical. High latency can lead to poor user experiences or even system failures in safety-critical applications.
Latency is highly dependent on three factors: the number of parameters in the model, the FLOPs required by the model, and hardware specific constraints. Because latency is extremely hardware dependent, it is often not a good choice for comparisons. However, it is the most crucial metric for real world applications.
During inference and training, a typical machine learning data path involves: (1) reading data from memory to the streaming multiprocessors on the GPU core (2) performing the required computations and (3) writing the results back to memory, which can happen asynchronously thanks to the GPU’s VRAM supprting async I/O. Hence, we have two possible bottlenecks: the computation that takes place on the GPU cores (\(T_{compute}\)), or the memory input/output (\(T_{compute}\)). The latency of an operation is therefore determined by the “slowest” of these two steps.
\[\text{latency} = \max(T_{memory}, T_{compute})\]If \(T_{memory} > T_{compute}\) we say the operation is memory-bound, else we say it is compute-bound. How to find out \(T_{memory}\) and \(T_{compute}\) for a specific neural network layer? \(T_{compute}\) is simply the amount of time the cuda cores need to perform the computations required. Assuming our layer performs a given amount of FLOPs, and the GPU can perform at most Floating Point Operations per Second, we have that
\[T_{compute} = \frac{layer FLOPs}{GPU FLOPS}\]Notice that the maximum GPU FLOPs depend on the precision we are operating at, so FLOPs at fp16 are not the same as FLOPs at int8.
The \(T_{compute}\) is determined by the amount of data we have to move from memory, usually High Bandwidth Memory (HBM) and the memory bandwidth. When executing a layer, we need to read from memory (a) the layer parameters and (b) the input activations. Therefore, we have that
\(T_{memory} = \frac{\text{size of model parameters} + \text{size of input activations} + \text{size of output activations}}{\text{memory bandwidth}}\) Again, the size of parameters is determined by the precision at which we are operating.
Throughput measures the number of inferences or predictions a model can make in a given time frame, usually one second. For an image classification task, these might be the number of images the model can classify in one second. While throughput is highly related to latency, they are not necessarily proportional. In order to increase throughput, we migh simply buy more GPUs and so get the chance to process more data in parallel, while the latency for a single input-output stays the same.
Energy consumption measures the energy required to perform an operation. This metric has gained increasing importance due to the environmental impact of large-scale machine learning and the operational costs associated with running models. Energy-efficient models not only reduce the carbon footprint but also lower the operational expenses for businesses deploying machine learning solutions. Techniques to minimize energy consumption include optimizing algorithms, utilizing energy-efficient hardware, and adopting more sustainable practices in data centers. Recent work has proposed models to estimate the carbon footprint for training models.
Additionally, energy is crucial for on-device training, which is necessary for all those scenarios where data cannot be shared due to privacy concerns. Unlike inference, training requires significant computational resources to adjust the model’s parameters. However, the computational and energy resources offered by edge devices are often scarce.
In general, VRAM memory access is the most energy consuming operation, requiring 100x more power than accessing SRAM and 200x than performing an ADD operation.
Peak activations refer to the maximum memory occupied by activations (outputs) produced by neurons in the network during the forward propagation process. At inference time they do not represent a significant bottleneck, because (a) the batch size it typically small and (b) we don’t need to store them for backpropagation. During training, the activations represent the most significant bottleneck, because we need to store them for gradient computation to update model weights.
It is important to point out that only activations of nonlinear layers introduce this problem. Assuming we have a layer \(a_{i+1} = \mathbf{w_i}^T a_i + b\). At some point during backprop, we need to compute the gradient of loss with respect to the activations \(a_i\). Applying the chain rule, we have that:
\[\frac{\partial \mathcal{L}}{\partial a_i} = \frac{\partial \mathcal{L}}{\partial a_{i+1}} \frac{\partial a_{i+1}}{\partial a_i} = \frac{\partial \mathcal{L}}{\partial a_{i+1}} \mathbf{w_i}^T\]if the layer is linear, whereas we have
\[\frac{\partial \mathcal{L}}{\partial a_i} = \frac{\partial \mathcal{L}}{\partial a_{i+1}} \frac{\partial a_{i+1}}{\partial a_i} = \frac{\partial \mathcal{L}}{\partial a_{i+1}} \mathbf{g}(a_i)\]where \(\mathbf{g}(a_i)\) depends on the activations at previous layer. This means that for a nonlinear layer, we need to store all the activations.
Model FLOPs utilization was proposed by Google as a metric to capture how well a specific model is using the hardware available. More specifically, it is defined as the ratio between the observed output throughput and the maximum theoretical throughtput.
\[\text{MFU} = \frac{Achieved FLOPs/s}{Theoretical Peak FLOPs/s}\]Importantly, the maximum theoretical throughput only accounts for FLOPs (forward and backward) and not for rematerialization + other computational overhead. A low MFU means that the model is underutilizing the hardware due to inefficiencies such as memory bottlenecks, poor parallelism, or excessive communication overhead.
This is an extremely important metric saying a lot about our implementation and training pipeline, as it captures the number of FLOPs a model theoretically utilizes compared to hardware peak FLOPs.
| Metric | explanation | primarily affects | primarily affected by | hardware independent |
|---|---|---|---|---|
| Parameters size | memory footprint of model | latency, storage | model architecture & implementation, floating point precision | ✅ |
| FLOPs ( and MACs) | number of operations required by model | latency, energy | model architecture & implementation | ✅ |
| Latency | time required for one inference | end user experience | Parameters size, FLOPs, Hardware | ❌ |
| Energy | power consumed by operation | training/inference cost | model architecture, implementation | ❌ |
| Throughput | inferences per time frame | end user experience, energy | all | ❌ |
| Peak activations | memory occupied by outputs | training memory consumption | all | ✅ |
After using different github libraries (see references), I found out that Pytorch actually has a built-in function to count flops.
I leave here a code snippet that shows how to compute the flops for a pytorch model only with forward or with forward and backward pass. We just need to provide the model and the input shapes for the model (or an input batch).
import torch
from torch.utils.flop_counter import FlopCounterMode
def get_flops(model, inp: Union[torch.Tensor, Tuple], with_backward=False):
istrain = model.training
model.eval()
inp = inp if isinstance(inp, torch.Tensor) else torch.randn(inp)
flop_counter = FlopCounterMode(mods=model, display=False, depth=None)
with flop_counter:
if with_backward:
model(inp).sum().backward()
else:
model(inp)
total_flops = flop_counter.get_total_flops()
if istrain:
model.train()
return total_flops
Say you want to use the snippet to compute the flops for a resnet18, the you would do something like the following.
from torchvision.models import resnet18
model = resnet18()
get_flops(model, (1, 3, 224, 224))
I leave the discussion about whther FLOPs are actually a good way of measuring efficiency to another blog post