<?xml version=”1.0” encoding=”utf-8”?>{% if page.xsl %}<?xml-stylesheet type=”text/xml” href=”{{ ‘/feed.xslt.xml’ absolute_url }}”?>{% endif %}<feed xmlns=”http://www.w3.org/2005/Atom” {% if site.lang %}xml:lang=”{{ site.lang }}”{% endif %}>Jekyll<link href=”{{ page.url absolute_url }}” rel=”self” type=”application/atom+xml” /><link href=”{{ ‘/’ absolute_url }}” rel=”alternate” type=”text/html” {% if site.lang %}hreflang=”{{ site.lang }}” {% endif %}/>{{ site.time date_to_xmlschema }}</updated>{{ page.url absolute_url xml_escape }}</id>{% assign title = site.title default: site.name %}{% if page.collection != “posts” %}{% assign collection = page.collection capitalize %}{% assign title = title append: “ ” append: collection %}{% endif %}{% if page.category %}{% assign category = page.category capitalize %}{% assign title = title append: “ ” append: category %}{% endif %}{% if title %}{{ title smartify xml_escape }}</title>{% endif %}{% if site.description %}{{ site.description xml_escape }}</subtitle>{% endif %}{% if site.author %}{{ site.author.name default: site.author xml_escape }}</name>{% if site.author.email %}{{ site.author.email xml_escape }}</email>{% endif %}{% if site.author.uri %}{{ site.author.uri xml_escape }}</uri>{% endif %}</author>{% endif %}{% if page.tags %}{% assign posts = site.tags[page.tags] %}{% else %}{% assign posts = site[page.collection] %}{% endif %}{% if page.category %}{% assign posts = posts where: “categories”, page.category %}{% endif %}{% unless site.show_drafts %}{% assign posts = posts where_exp: “post”, “post.draft != true” %}{% endunless %}{% assign posts = posts sort: “date” reverse %}{% assign posts_limit = site.feed.posts_limit default: 10 %}{% for post in posts limit: posts_limit %}<entry{% if post.lang %}{{“ “}}xml:lang=”{{ post.lang }}”{% endif %}>{% assign post_title = post.title smartify strip_html normalize_whitespace xml_escape %}{{ post_title }}<link href=”{{ post.url absolute_url }}” rel=”alternate” type=”text/html” title=”{{ post_title }}” />{{ post.date date_to_xmlschema }}</published>{{ post.last_modified_at default: post.date date_to_xmlschema }}</updated>{{ post.id absolute_url xml_escape }}</id>{% assign excerpt_only = post.feed.excerpt_only default: site.feed.excerpt_only %}{% unless excerpt_only %}<content type=”html” xml:base=”{{ post.url absolute_url xml_escape }}”><![CDATA[{{ post.content strip }}]]></content>{% endunless %}{% assign post_author = post.author default: post.authors[0] default: site.author %}{% assign post_author = site.data.authors[post_author] default: post_author %}{% assign post_author_email = post_author.email default: nil %}{% assign post_author_uri = post_author.uri default: nil %}{% assign post_author_name = post_author.name default: post_author %}{{ post_author_name default: “” xml_escape }}</name>{% if post_author_email %}{{ post_author_email xml_escape }}</email>{% endif %}{% if post_author_uri %}{{ post_author_uri xml_escape }}</uri>{% endif %}</author>{% if post.category %}<category term=”{{ post.category xml_escape }}” />{% elsif post.categories %}{% for category in post.categories %}<category term=”{{ category xml_escape }}” />{% endfor %}{% endif %}{% for tag in post.tags %}<category term=”{{ tag xml_escape }}” />{% endfor %}{% assign post_summary = post.description default: post.excerpt %}{% if post_summary and post_summary != empty %}<summary type="html"><![CDATA[{{ post_summary strip_html normalize_whitespace }}]]></summary>{% endif %}{% assign post_image = post.image.path default: post.image %}{% if post_image %}{% unless post_image contains “://” %}{% assign post_image = post_image absolute_url %}{% endunless %}<media:thumbnail xmlns:media=”http://search.yahoo.com/mrss/” url=”{{ post_image xml_escape }}” /><media:content medium=”image” url=”{{ post_image xml_escape }}” xmlns:media=”http://search.yahoo.com/mrss/” />{% endif %}</entry>{% endfor %}</feed>

Fine-Tuning LLaMA 3.2-11B-Vision for Product Descriptions

9 minute read

Published: February 23, 2025

Fine-Tuning LLaMA 3.2-11B-Vision for Product Descriptions The complete implementation is available on GitHub.Introduction:Large vision-language models (LVLMs) like LLaMA 3.2-11B-Vision have revolutionized AI-powered content generation. However, their performance can be significantly enhanced when fine-tuned on domain-specific data. In my recent project, I fine-tuned LLaMA 3.2-11B-Vision to generate high-quality product descriptions for images of girls’ clothing.The goal was to create engaging, detailed, and accurate descriptions that could be used for e-commerce platforms, enhancing product listings and improving customer experience. By leveraging a carefully curated dataset and optimizing the fine-tuning process, I was able to achieve impressive results. In this blog post, I’ll walk through the entire process—dataset preparation, fine-tuning methodology, challenges faced, and key takeaways.Stay tuned to learn how you can fine-tune LLaMA 3.2-11B-Vision for your own vision-language tasks!SetupJust open the colab version:or clone the repogit clone https://github.com/ramintoosi/product_descriptioncd product_descriptionand follow the unsloth installation guide.Dataset PreparationBefore fine-tuning LLaMA 3.2-11B-Vision, the first step is to acquire and prepare the dataset. The dataset consists of images of girls’ clothing along with corresponding product descriptions, which will be used to train the model to generate high-quality text based on image inputs. I gathered this dataset by crawling a web store, ensuring that it contains diverse and well-structured product listings. While I cannot share the crawling code, the dataset itself is publicly available.To get started, we need to download the dataset from Google Drive and extract it into a working directory. The following commands will accomplish this:# Install gdown if not already installedpip install gdown # Download the dataset from Google Drivegdown --id 14PptNxqI7D6YuTiPOjt1H0uLaa8Qr0qF # Unzip the dataset into the 'data' directoryunzip -q product_description_data.zip -d data Once extracted, the dataset will be available in the data/ directory, ready for preprocessing. In the next section, we will explore the dataset structure and prepare it for fine-tuning.Loading and Converting the DataWith the dataset downloaded and extracted, the next step is to load the data, clean it, and convert it into a format suitable for fine-tuning LLaMA 3.2-11B-Vision with Unsloth. This involves: Loading the dataset from a CSV file that contains product names, brand names, and image paths. Cleaning the data by removing unnecessary codes, model numbers, and redundant information from product names. Converting the data into a structured conversation format that aligns with the input-output style expected by LLaMA 3.2-11B-Vision.Step 1: Loading and Cleaning the DataThe following Python function reads the dataset, drops irrelevant columns, removes duplicate entries, and applies text cleaning rules to refine product names:import osfrom typing import TypedDictimport pandas as pdclass Data(TypedDict): name: str brand: str image_path: strdef load_data(data_root: str) -> list[Data]: """ Load data from csv file :param data_root: data root folder where data.csv and images are stored :return: list of dictionaries, keys are column names and values are data """ data = pd.read_csv(os.path.join(data_root, 'data.csv')) data.drop(columns=['site_category', 'supply_category'], inplace=True) clean_data(data) return data.to_dict(orient='records')def clean_data(data: pd.DataFrame): """ Clean data by removing product codes and model numbers. :param data: data frame """ # Remove duplicate rows data.drop_duplicates(inplace=True) # Remove codes, model numbers, and specific patterns pattern = r'\b(کد|مدل)\b(\s+[A-Za-z0-9]+)?|\bمجموعه\b\s+(\d+)\s+(\w+)' data["name"] = data["name"].str.replace(pattern, '', regex=True) # Remove any remaining standalone alphanumeric codes in English pattern2 = r'\b[A-Za-z0-9_-]+\b' data["name"] = data["name"].str.replace(pattern2, '', regex=True)data = load_data('./data')print(data[0]) # Sample outputSample Output:{ "name": "ست تی شرت آستین بلند و شلوار بچگانه سپیدپوش ماشین پلیس", "brand": "سپیدپوش", "image_path": "data/dataset/image/59e8b029-864c-4788-bf4f-25d9f0ad494d.jpg"}Step 2: Converting Data to a Conversation FormatLLaMA 3.2-11B-Vision expects data in a structured conversational format, where the user provides an instruction along with an image, and the model generates a response. The following function transforms each data sample into this format:from PIL import Imageinstruction = """Create a Short Product description based on the provided ##PRODUCT BRAND NAME## and the image.Only return description. The description should be SEO optimized and for a better mobile search experience.##PRODUCT BRAND NAME##: {brand_name}"""def convert_to_conversation(sample): image_path = sample["image_path"].replace('dataset/', '') conversation = [ { "role": "user", "content" : [ {"type" : "text", "text" : instruction.format(brand_name=sample["brand"])}, {"type" : "image", "image_url" : f'file://{image_path}'} ] }, { "role" : "assistant", "content" : [ {"type" : "text", "text" : sample["name"]} ] }, ] return { "messages" : conversation }converted_dataset = [convert_to_conversation(sample) for sample in data]print(converted_dataset[0]) # Sample outputSample Output:{ "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Create a Short Product description based on the provided ##PRODUCT BRAND NAME## and the image.\nOnly return description. The description should be SEO optimized and for a better mobile search experience.\n\n##PRODUCT BRAND NAME##: سپیدپوش\n" }, { "type": "image", "image_url": "file://data/image/59e8b029-864c-4788-bf4f-25d9f0ad494d.jpg" } ] }, { "role": "assistant", "content": [ { "type": "text", "text": "ست تی شرت آستین بلند و شلوار بچگانه سپیدپوش ماشین پلیس" } ] } ]}Efficient Image Handling: Avoiding Memory OverloadIn the original Unsloth library tutorial, image inputs are loaded into memory using:{"type": "image", "image": Image(image_path)}This method loads all images at once, which can cause memory crashes when dealing with large datasets like ours.To avoid excessive memory usage, we store image file paths instead of loading images into memory. This way, the model loads images dynamically during training instead of keeping them in RAM.{"type": "image", "image_url": f'file://{image_path}'}Loading and Fine-tuning the LLaMA 3.2-11B-Vision Model with UnslothNow that we have our dataset ready, the next step is to load and configure the model for fine-tuning. We will use Unsloth’s FastVisionModel, which provides optimized loading and memory-efficient training.Step 1: Import Required Librariesfrom unsloth import FastVisionModel # FastLanguageModel for LLMsimport torchStep 2: Load Pretrained Model and TokenizerWe initialize the LLaMA 3.2-11B-Vision-Instruct model with 4-bit quantization to optimize memory usage.model, tokenizer = FastVisionModel.from_pretrained( "unsloth/Llama-3.2-11B-Vision-Instruct", load_in_4bit=True, # Use 4-bit quantization to reduce memory usage use_gradient_checkpointing="unsloth", # Activates checkpointing for long context)✅ 4-bit quantization reduces memory usage significantly, allowing us to fine-tune on consumer-grade GPUs.✅ Gradient checkpointing helps handle long-context sequences efficiently.Checking the Model’s Pretrained Vision CapabilitiesBefore fine-tuning, it’s useful to check whether LLaMA 3.2-11B-Vision already understands and analyzes images effectively. We do this by running an inference test using the pretrained model.Step 2.1: Enable Inference ModeFastVisionModel.for_inference(model) # Switch to inference modeStep 2.2: Prepare Image Inputfrom PIL import Imageimage = Image.open(data[0]["image_path"].replace('dataset/', ''))messages = [ {"role": "user", "content": [ {"type": "image"}, {"type": "text", "text": instruction} ]}]Step 2.3: Tokenize Input for the Modelinput_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)inputs = tokenizer( image, input_text, add_special_tokens=False, return_tensors="pt",).to("cuda")Step 2.4: Generate the Model’s Responsefrom transformers import TextStreamertext_streamer = TextStreamer(tokenizer, skip_prompt=True)_ = model.generate(**inputs, streamer=text_streamer, max_new_tokens=128, use_cache=True, temperature=1.5, min_p=0.1)Model-generated description: “The image showcases a children’s pajama set from the brand {brand_name}. The pajama shirt is a long-sleeved grey and white striped shirt featuring a playful police car design on the front. The police car is depicted in blue, with a white dome on top, adorned with a red siren light, and sporting black wheels. Below the car, the word ‘POLICE’ is written in blue text. The pajama bottoms are solid black, made from a stretchy fabric designed to move with the wearer.”Original dataset caption: “ست تی شرت آستین بلند و شلوار بچگانه سپیدپوش ماشین پلیس”Step 3: Enable Parameter-Efficient Fine-Tuning (PEFT)To fine-tune the model efficiently, we use LoRA (Low-Rank Adaptation). This method only trains specific layers, drastically reducing GPU memory consumption.model = FastVisionModel.get_peft_model( model, finetune_vision_layers=False, # Vision layers are frozen to focus on text generation finetune_language_layers=True, # Enable fine-tuning of language layers finetune_attention_modules=True, # Fine-tune attention layers for better adaptation finetune_mlp_modules=True, # Fine-tune MLP layers for better generalization r=16, # Controls rank of LoRA adaptation; higher values improve accuracy but increase overfitting risk lora_alpha=16, # LoRA scaling factor (recommended: equal to `r`) lora_dropout=0, # No dropout for stable fine-tuning bias="none", # No additional bias parameters random_state=3407, # Ensures reproducibility use_rslora=False, # Rank-stabilized LoRA disabled (can improve LoRA stability in some cases) loftq_config=None, # LoftQ disabled (for further quantization efficiency) # target_modules="all-linear", # Optional: Specifies which layers to adapt)Step 4: Fine-Tuning the ModelNow that we’ve confirmed the pretrained model can analyze images, we move on to fine-tuning it for generating concise, SEO-optimized product descriptions. We use the Trainer API to fine-tune the model on our dataset.trainer_stats = trainer.train()Evaluating the Fine-Tuned ModelNow that our model is trained, let’s test it on a sample image and compare its new description with the original dataset caption.Step 1: Select a Sample & Enable Inference Modesample_idx = 157FastVisionModel.for_inference(model) # Switch back to inference modeStep 2: Load and Prepare the Imagefrom PIL import Imageimage = Image.open(data[sample_idx]["image_path"].replace('dataset/', ''))messages = [ {"role": "user", "content": [ {"type": "image"}, {"type": "text", "text": instruction.format(brand_name=data[sample_idx]["brand"])} ]}]Step 3: Tokenize Input for the Modelinput_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)inputs = tokenizer( image, input_text, add_special_tokens=False, return_tensors="pt",).to("cuda")Step 4: Generate a Descriptionfrom transformers import TextStreamertext_streamer = TextStreamer(tokenizer, skip_prompt=True)_ = model.generate(**inputs, streamer=text_streamer, max_new_tokens=128, use_cache=True, temperature=1.5, min_p=0.1)ResultsModel Caption: جوراب ساق بلند دخترانه کاتامیناOriginal Caption: جوراب دخترانه کاتامیناGitHub RepositoryThe complete implementation is available on GitHub and .Reference[1] Unsloth Tutorials[2] ChatGPT

Diffusion From Scratch To Generate Cute Anime Faces!

14 minute read

Published: November 27, 2024

Diffusion From Scratch To Generate Cute Anime Faces! The complete implementation is available on GitHub. The complete implementation is available on GitHub.Introduction:Diffusion models have revolutionized the field of AI art, enabling the generation of stunning visuals, including adorable anime faces. While prebuilt tools like Stable Diffusion make it easy to create such images, building a diffusion model from scratch offers invaluable benefits. It deepens your understanding of the underlying principles, gives you the flexibility to customize and innovate, and provides a hands-on learning experience that can be both rewarding and empowering. In this blog post, we’ll explore how to create a diffusion model step-by-step, equipping you with the knowledge to generate charming anime faces and unlock your creative potential.SetupClone the Repositorygit clone https://github.com/ramintoosi/diffusion-from-scratchinstall requirementspip install -r requirements.txttrain the modelpython main.py traininferencepython main.py inferenceOk, let’s go through the details.Diffusion ModelDiffusion models operate on a simple yet powerful intuition: they learn to generate data by reversing a gradual noise process. Imagine starting with a clear image and progressively adding random noise to it, step by step, until it becomes completely unrecognizable. Diffusion models essentially learn this “noising” process in reverse—they begin with random noise and iteratively refine it to recover the original, clear image. By training on vast datasets, they master this denoising process, enabling them to generate entirely new images from pure noise, guided by patterns and structures learned during training. This step-by-step refinement makes diffusion models particularly adept at producing high-quality and detailed outputs.Model Architecture ImplementationForward PassThe DiffusionForwardProcess` class implements the forward diffusion process, which systematically adds noise to an image over a series of time steps. This gradual “noising” process prepares the model for learning to reverse it during training.class DiffusionForwardProcess: """ Implements the forward process of the diffusion model. """ def __init__(self, num_time_steps: int =1000, beta_start: float = 1e-4, beta_end: float = 0.02 ): """ Initializes the DiffusionForwardProcess with the given parameters. :param num_time_steps: Number of time steps in the diffusion process. :param beta_start: Starting value of beta. :param beta_end: Ending value of beta. """ self.betas = torch.linspace(beta_start, beta_end, num_time_steps) self.alphas = 1 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) self.sqrt_alpha_bars = torch.sqrt(self.alpha_bars) self.sqrt_one_minus_alpha_bars = torch.sqrt(1 - self.alpha_bars) def add_noise(self, original: Tensor, noise: Tensor, t: Tensor) -> Tensor: """ Adds noise to the original image at the given time step t. :param original: Input Image :param noise: Random Noise Tensor sampled from Normal Dist :param t: timestep :return: Noisy image tensor """ sqrt_alpha_bar_t = self.sqrt_alpha_bars.to(original.device)[t] sqrt_one_minus_alpha_bar_t = self.sqrt_one_minus_alpha_bars.to(original.device)[t] # Broadcast to multiply with the original image. sqrt_alpha_bar_t = sqrt_alpha_bar_t[:, None, None, None] sqrt_one_minus_alpha_bar_t = sqrt_one_minus_alpha_bar_t[:, None, None, None] # Return return (sqrt_alpha_bar_t * original) \ + \ (sqrt_one_minus_alpha_bar_t * noise)Reverse PassThe DiffusionReverseProcess class implements the reverse denoising process, which is the core of the generative aspect of diffusion models. Starting with pure noise, this process iteratively predicts and reconstructs cleaner versions of the image until the original-like data emerges.class DiffusionReverseProcess: """ Implements the reverse process of the diffusion model. """ def __init__(self, num_time_steps: int = 1000, beta_start: float = 1e-4, beta_end: float = 0.02 ): """ Initializes the DiffusionReverseProcess with the given parameters. :param num_time_steps: Number of time steps in the diffusion process. :param beta_start: Starting value of beta. :param beta_end: Ending value of beta. """ # Precomputing beta, alpha, and alpha_bar for all t's. self.b = torch.linspace(beta_start, beta_end, num_time_steps) # b -> beta self.a = 1 - self.b # a -> alpha self.a_bar = torch.cumprod(self.a, dim=0) # a_bar = alpha_bar def sample_prev_timestep(self, xt: Tensor, noise_pred: Tensor, t) -> (Tensor, Tensor): """ Samples the previous timestep image given the current timestep image and noise prediction. :param xt: Image tensor at timestep t of shape -> B x C x H x W :param noise_pred: Noise tensor predicted by the model at timestep t of shape -> B x C x H x W :param t: timestep :return: predicted x_t-1 and x0 """ # Original Image Prediction at timestep t x0 = xt - (torch.sqrt(1 - self.a_bar.to(xt.device)[t]) * noise_pred) x0 = x0 / torch.sqrt(self.a_bar.to(xt.device)[t]) x0 = torch.clamp(x0, -1., 1.) # mean of x_(t-1) mean = (xt - ((1 - self.a.to(xt.device)[t]) * noise_pred) / (torch.sqrt(1 - self.a_bar.to(xt.device)[t]))) mean = mean / (torch.sqrt(self.a.to(xt.device)[t])) # only return mean if t == 0: return mean, x0 else: variance = (1 - self.a_bar.to(xt.device)[t - 1]) / (1 - self.a_bar.to(xt.device)[t]) variance = variance * self.b.to(xt.device)[t] sigma = variance ** 0.5 z = torch.randn(xt.shape).to(xt.device) return mean + sigma * z, x0Time EmbeddingThe get_time_embedding function generates a time-step embedding, transforming scalar time-step values into high-dimensional vector representations. These embeddings are critical for encoding temporal information that guides the diffusion model in processing noise at specific timesteps.def get_time_embedding(time_steps: torch.Tensor, t_emb_dim: int) -> torch.Tensor: """ Transform a scalar time-step into a vector representation of size t_emb_dim. :param time_steps: 1D tensor of size -> (Batch, ) :param t_emb_dim: Embedding Dimension -> for ex: 128 (scalar value) :return tensor of size -> (B, t_emb_dim) """ assert t_emb_dim % 2 == 0, "time embedding must be divisible by 2." factor = 2 * torch.arange(start=0, end=t_emb_dim // 2, dtype=torch.float32, device=time_steps.device ) / t_emb_dim factor = 10000 ** factor t_emb = time_steps[:, None] # B -> (B, 1) t_emb = t_emb / factor # (B, 1) -> (B, t_emb_dim//2) t_emb = torch.cat([torch.sin(t_emb), torch.cos(t_emb)], dim=1) # (B , t_emb_dim) return t_embThe TimeEmbedding class is a simple neural network module that transforms a time-step embedding into the desired output dimension. This transformation is typically used to align the time embedding with the dimensions required for downstream operations in the diffusion model.class TimeEmbedding(nn.Module): """ Maps the Time Embedding to the Required output Dimension. """ def __init__(self, n_out: int, # Output Dimension t_emb_dim: int = 128 # Time Embedding Dimension ): super().__init__() # Time Embedding Block self.te_block = nn.Sequential( nn.SiLU(), nn.Linear(t_emb_dim, n_out) ) def forward(self, x): return self.te_block(x)Conv BlockThe NormActConv class is a module that sequentially applies Group Normalization, Activation, and Convolution operations. This modular design simplifies the construction of neural networks used in diffusion models, particularly for processing image data.class NormActConv(nn.Module): """ Perform GroupNorm, Activation, and Convolution operations. """ def __init__(self, in_channels: int, out_channels: int, num_groups: int = 8, kernel_size: int = 3, norm: bool = True, act: bool = True ): super().__init__() # GroupNorm self.g_norm = nn.GroupNorm( num_groups, in_channels ) if norm is True else nn.Identity() # Activation self.act = nn.SiLU() if act is True else nn.Identity() # Convolution self.conv = nn.Conv2d( in_channels, out_channels, kernel_size, padding=(kernel_size - 1) // 2 ) def forward(self, x): x = self.g_norm(x) x = self.act(x) x = self.conv(x) return xSelf AttentionThe SelfAttentionBlock class is a neural network module that applies Group Normalization followed by Multi-Headed Self-Attention to capture long-range dependencies in the input data. This is particularly useful in processing spatial information in images for tasks like denoising and generation.class SelfAttentionBlock(nn.Module): """ Perform GroupNorm and Multi-headed Self Attention operation. """ ...Downsample and Upsample BlocksThe Downsample class performs downsampling operations on input tensors, reducing their spatial dimensions (height and width) by a factor of k. It provides two methods for downsampling—Convolution-based and Max-pooling-based, with the option to combine both for enhanced feature representation. Upsampling module follows a similar design pattern, but in the reverse direction, increasing the spatial dimensions of the input tensor.class Downsample(nn.Module): """ Perform Down sampling by the factor of k across Height and Width. """ def __init__(self, in_channels: int, out_channels: int, k: int = 2, # Downsampling factor use_conv: bool = True, # If Downsampling using conv-block use_mpool: bool = True # If Downsampling using max-pool ): super(Downsample, self).__init__() self.use_conv = use_conv self.use_mpool = use_mpool # Downsampling using Convolution self.cv = nn.Sequential( nn.Conv2d(in_channels, in_channels, kernel_size=1), nn.Conv2d( in_channels, out_channels // 2 if use_mpool else out_channels, kernel_size=4, stride=k, padding=1 ) ) if use_conv else nn.Identity() # Downsampling using Maxpool self.mpool = nn.Sequential( nn.MaxPool2d(k, k), nn.Conv2d( in_channels, out_channels // 2 if use_conv else out_channels, kernel_size=1, stride=1, padding=0 ) ) if use_mpool else nn.Identity() def forward(self, x): if not self.use_conv: return self.mpool(x) if not self.use_mpool: return self.cv(x) return torch.cat([self.cv(x), self.mpool(x)], dim=1)Unet PartsThe DownC, MidC, and UpC modules are key components of a U-Net-like architecture tailored for image generation tasks, such as diffusion models. These modules work together to process and transform input data through hierarchical downsampling, mid-level refinement, and upsampling.class DownC(nn.Module): """ Perform Down-convolution on the input using following approach. 1. Conv + TimeEmbedding 2. Conv 3. Skip-connection from input x. 4. Self-Attention 5. Skip-Connection from 3. 6. Downsampling """class MidC(nn.Module): """ Refine the features obtained from the DownC block. It refines the features using following operations: 1. Resnet Block with Time Embedding 2. A Series of Self-Attention + Resnet Block with Time-Embedding """class UpC(nn.Module): """ Perform Up-convolution on the input using following approach. 1. Upsampling 2. Conv + TimeEmbedding 3. Conv 4. Skip-connection from 1. 5. Self-Attention 6. Skip-Connection from 3. """1. DownC: Down-Convolution BlockResponsible for extracting multi-scale features while reducing spatial resolution. It captures hierarchical features and incorporates temporal context via time embeddings.Key Steps: Convolution + Time Embedding: Extracts spatial features while integrating temporal information (important for diffusion models). Additional Convolution: Deepens feature extraction. Skip-Connection: Creates a residual path from the original input, preserving critical low-level details. Self-Attention: Enables the model to capture long-range dependencies in spatial data. Skip-Connection: Adds the features from step 3 back into the flow, enhancing gradient flow and feature reuse. Downsampling: Reduces spatial resolution, allowing deeper layers to focus on abstract patterns.2. MidC: Middle BlockServes as the bottleneck in the U-Net, refining features extracted by DownC while integrating global context.Key Steps: ResNet Block with Time Embedding: Refines features using residual connections while incorporating temporal embeddings for time-awareness. Self-Attention + ResNet Block: A sequence of self-attention layers and ResNet blocks further enriches features by combining spatial attention with robust feature refinement.3. UpC: Up-Convolution BlockReconstructs the spatial dimensions by progressively upsampling while merging features from earlier layers via skip-connections.Key Steps: Upsampling: Increases spatial resolution using learned or interpolation-based methods. Convolution + Time Embedding: Processes upscaled features while integrating temporal information. Additional Convolution: Refines upsampled features. Skip-Connection: Reintroduces earlier features, ensuring high-resolution details are preserved. Self-Attention: Captures global spatial dependencies, essential for generating coherent outputs. Skip-Connection: Combines processed features with those from step 3 for enhanced reconstruction.Integration in the U-Net: Encoder (DownC): Captures features at multiple resolutions, gradually reducing spatial size while increasing feature richness. Bottleneck (MidC): Acts as a bridge, blending abstract, low-resolution features with global dependencies. Decoder (UpC): Reconstructs the image, integrating skip-connections from the encoder to ensure high-resolution details are preserved.class Unet(nn.Module): """ U-net architecture """ def __init__(self, im_channels: int = 1, # RGB down_ch=None, mid_ch=None, up_ch=None, down_sample=None, t_emb_dim: int = 128, num_downc_layers: int = 2, num_midc_layers: int = 2, num_upc_layers: int = 2 ): super(Unet, self).__init__() if down_sample is None: down_sample = [True, True, False] if up_ch is None: up_ch = [256, 128, 64, 16] if mid_ch is None: mid_ch = [256, 256, 128] if down_ch is None: down_ch = [32, 64, 128, 256] self.im_channels = im_channels self.down_ch = down_ch self.mid_ch = mid_ch self.up_ch = up_ch self.t_emb_dim = t_emb_dim self.down_sample = down_sample self.num_downc_layers = num_downc_layers self.num_midc_layers = num_midc_layers self.num_upc_layers = num_upc_layers self.up_sample = list(reversed(self.down_sample)) # [False, True, True] # Initial Convolution self.cv1 = nn.Conv2d(self.im_channels, self.down_ch[0], kernel_size=3, padding=1) # Initial Time Embedding Projection self.t_proj = nn.Sequential( nn.Linear(self.t_emb_dim, self.t_emb_dim), nn.SiLU(), nn.Linear(self.t_emb_dim, self.t_emb_dim) ) # DownC Blocks self.downs = nn.ModuleList([ DownC( self.down_ch[i], self.down_ch[i + 1], self.t_emb_dim, self.num_downc_layers, self.down_sample[i] ) for i in range(len(self.down_ch) - 1) ]) # MidC Block self.mids = nn.ModuleList([ MidC( self.mid_ch[i], self.mid_ch[i + 1], self.t_emb_dim, self.num_midc_layers ) for i in range(len(self.mid_ch) - 1) ]) # UpC Block self.ups = nn.ModuleList([ UpC( self.up_ch[i], self.up_ch[i + 1], self.t_emb_dim, self.num_upc_layers, self.up_sample[i] ) for i in range(len(self.up_ch) - 1) ]) # Final Convolution self.cv2 = nn.Sequential( nn.GroupNorm(8, self.up_ch[-1]), nn.Conv2d(self.up_ch[-1], self.im_channels, kernel_size=3, padding=1) ) def forward(self, x, t): out = self.cv1(x) # Time Projection t_emb = get_time_embedding(t, self.t_emb_dim) t_emb = self.t_proj(t_emb) # DownC outputs down_outs = [] for down in self.downs: down_outs.append(out) out = down(out, t_emb) # MidC outputs for mid in self.mids: out = mid(out, t_emb) # UpC Blocks for up in self.ups: down_out = down_outs.pop() out = up(out, down_out, t_emb) # Final Conv out = self.cv2(out) return outDataDownload the Anime Face Dataset and put it in ./data/anime.TrainingThe training function is a straight forward PyTorch training loop that iterates over the dataset and updates the model parameters using the Adam optimizer. The loss function is the Mean Squared Error (MSE) loss, which measures the difference between the predicted and target images.InferenceUsing the inference module, one can generate images using the trained model. The inference process involves sampling noise tensors and iteratively predicting the previous timestep image until the original-like image is reconstructed.def generate(cfg: CONFIG) -> Tensor: """ Generate Image using trained model. :param cfg: config :return: image tensor """ # Device device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu') # print(f'Device: {device}\n') # Initialize Diffusion Reverse Process drp = DiffusionReverseProcess() # Set model to eval mode model = torch.load(cfg.model_path).to(device) model.eval() # Generate Noise sample from N(0, 1) xt = torch.randn(1, cfg.in_channels, cfg.img_size, cfg.img_size).to(device) # Denoise step by step by going backward. with torch.no_grad(): for t in reversed(range(cfg.num_timesteps)): noise_pred = model(xt, torch.as_tensor(t).unsqueeze(0).to(device)) xt, x0 = drp.sample_prev_timestep(xt, noise_pred, torch.as_tensor(t).to(device)) # Convert the image to proper scale xt = torch.clamp(xt, -1., 1.).detach().cpu() xt = (xt + 1) / 2 return xtResultsGitHub RepositoryThe complete implementation is available on GitHub.Reference[1] DDPM FROM SCRATCH[2] ChatGPT

Running Llama 3.2 in Rust

15 minute read

Published: October 28, 2024

Running Llama 3.2 in Rust The complete implementation is available on GitHub. The complete implementation is available on GitHub.Introduction:With the rapid growth of AI and natural language processing, efficient language model deployment has become a key focus for developers. Models like Llama 3.2, known for their performance and flexibility, open doors for sophisticated applications in text generation, chatbots, summarization, and more. But, running these models at high speeds—especially on GPUs—requires a language that balances performance with control.In this post, we’ll explore how to set up and run Llama 3.2 in Rust, a language gaining popularity for its system-level access, memory safety, and concurrency features. Using the llama.cpp library as our backend, we’ll implement a flexible language model interface with Rust.SetupTo get started with Llama 3.2 in Rust, we’ll first clone the project repository and set up the dependencies using Cargo, Rust’s package manager. Since this project depends on the llama.cpp backend, we need to clone it recursively to ensure all submodules are included. Follow these steps to prepare the environment:Clone the RepositoryFirst, clone the repository recursively to pull in all necessary submodules:git clone --recursive https://github.com/ramintoosi/llama-rustcd llama-rust Note: The --recursive flag is essential for including the llama.cpp bindings, which provides the backend support for model inference.Cargo DependenciesThe project’s Cargo.toml file includes the following dependencies:[dependencies]llama-cpp-2 = { path = "llama-cpp-rs/llama-cpp-2", features = ["cuda"] }hf-hub = "0.3.2"clap = { version = "4.5.19", features = ["derive"] }anyhow = "1.0.89"encoding_rs = "0.8.34"log = "0.4.22"[features]cuda = ["llama-cpp-2/cuda"] llama-cpp-2: Connects the project to the llama.cpp Rust binding, with an optional cuda feature for GPU support.Arguments ParsingThe Args module, built with the clap library, efficiently manages the command-line arguments required for model selection, configuration, and customization. Here’s a breakdown of how each component works, focusing on the design and functionality of the different parameters.use anyhow::{anyhow, Context};use std::path::PathBuf;use clap::{Parser, Subcommand};use hf_hub::api::sync::ApiBuilder;use llama_cpp_2::model::params::kv_overrides::ParamOverrideValue;use std::str::FromStr;#[derive(Subcommand, Debug, Clone)]pub enum Model { /// Use an already downloaded model #[clap(name = "local")] Local { /// The path to the model. e.g. `../hub/models--TheBloke--Llama-2-7B-Chat-GGUF/blobs/08a5566d61d7cb6b420c3e4387a39e0078e1f2fe5f055f3a03887385304d4bfa` /// or `./llama-3.2-1b-instruct-q8_0.gguf` path: PathBuf, }, /// Download a model from huggingface (or use a cached version) #[clap(name = "hf-model")] HuggingFace { /// the repo containing the model. e.g. `TheBloke/Llama-2-7B-Chat-GGUF` repo: String, /// the model name. e.g. `llama-2-7b-chat.Q4_K_M.gguf` model: String, },}impl Model { /// Convert the model to a path - may download from huggingface pub fn get_or_load(self) -> anyhow::Result<PathBuf> { match self { Model::Local { path } => Ok(path), Model::HuggingFace { model, repo } => ApiBuilder::new() .with_progress(true) .build() .with_context(|| "unable to create huggingface api")? .model(repo) .get(&model) .with_context(|| "unable to download model"), } }}#[derive(clap::ValueEnum, Clone, Debug)]pub enum Mode { Chat, Completion,}#[derive(Parser, Debug, Clone)]pub struct Args { /// The path to the model #[command(subcommand)] pub model: Model, /// The mode of the code: completion or chat #[clap(value_enum, short = 'm', long, default_value = "chat")] pub mode: Mode, // /// The prompt to use - valid only if the mode is `completion` // #[clap(short = 'p', long, required_if_eq("mode", "completion"))] // prompt: Option<String>, /// set the length of the prompt + output in tokens #[clap(long, default_value_t = 512)] pub max_token: u32, /// override some parameters of the model #[clap(short = 'o', value_parser = parse_key_val)] pub key_value_overrides: Vec<(String, ParamOverrideValue)>, /// how many layers to keep on the gpu - zero is cpu mode #[clap( short = 'g', long, help = "how many layers to keep on the gpu - zero is cpu mode (default: 0)" )] pub n_gpu_layers: u32, /// set the seed for the RNG #[clap(short = 's', long, default_value_t=561371)] pub seed: u32, /// number of threads to use during generation #[clap( long, help = "number of threads to use during generation (default: use all available threads)" )] pub threads: Option<i32>, #[clap( long, help = "number of threads to use during batch and prompt processing (default: use all available threads)" )] pub threads_batch: Option<i32>, // /// size of the prompt context // #[clap( // short = 'c', // long, // help = "size of the prompt context (default: loaded from the model)" // )] // pub ctx_size: Option<NonZeroU32>, /// show the token/s speed at the end of each turn #[clap(short = 'v', long, action)] pub verbose: bool,}/// Parse a single key-value pairfn parse_key_val(s: &str) -> anyhow::Result<(String, ParamOverrideValue)> { let pos = s .find('=') .ok_or_else(|| anyhow!("invalid KEY=value: no `=` found in `{}`", s))?; let key = s[..pos].parse()?; let value: String = s[pos + 1..].parse()?; let value = i64::from_str(&value) .map(ParamOverrideValue::Int) .or_else(|_| f64::from_str(&value).map(ParamOverrideValue::Float)) .or_else(|_| bool::from_str(&value).map(ParamOverrideValue::Bool)) .map_err(|_| anyhow!("must be one of i64, f64, or bool"))?; Ok((key, value))}Model Enum for Model SelectionThe Model enum defines two ways to specify a language model: Local Model (local): This variant lets users provide a path to a locally downloaded model file. Hugging Face Model (hf-model): This variant enables automatic downloading of a model from Hugging Face’s repositories. It takes in the repo and model names and uses hf_hub to fetch the specified model, caching it for future use. The get_or_load method in Model abstracts the logic for loading models, using Hugging Face’s API if needed.Mode Enum for Operation ModeThe Mode enum controls the inference mode: Chat: Enables a conversational interaction with the model, allowing for an ongoing dialogue. Completion: Generates text completions based on an initial prompt.The mode argument defaults to Chat, but users can specify either mode by passing -m completion or -m chat.The Args StructThe Args struct organizes all of the command-line arguments into a clean and accessible structure. Let’s look at each parameter: model: A subcommand that lets the user specify a model, either by path or by Hugging Face repository, as described in the Model enum. mode: Specifies the operational mode, either chat or completion. max_token: Sets the maximum length for prompt and output tokens. The default value is 512, allowing some control over the model’s generation length. key_value_overrides: Provides flexibility by allowing the user to pass key-value overrides for specific model parameters. n_gpu_layers: Specifies the number of layers to run on the GPU, with a default of 0, which uses the CPU. seed: Sets a seed for the random number generator, allowing for reproducible output across runs. This defaults to 561371. threads and threads_batch: These arguments allow the user to fine-tune performance by specifying the number of threads used for generation and batch processing, respectively. By default, it uses all available threads. verbose: If set, this flag displays token processing speed after each generation, useful for performance monitoring. Llama 3.2 InferenceThe LLM struct is the central component for handling Llama model inference, encapsulating essential components like the model, backend, and configuration parameters.use std::ffi::CString;use std::io::Write;use std::num::NonZeroU32;use std::pin::pin;use llama_cpp_2::context::LlamaContext;use llama_cpp_2::context::params::LlamaContextParams;use llama_cpp_2::llama_backend::LlamaBackend;use llama_cpp_2::llama_batch::LlamaBatch;use llama_cpp_2::model::{AddBos, LlamaModel, Special};use llama_cpp_2::model::params::LlamaModelParams;use llama_cpp_2::token::data_array::LlamaTokenDataArray;use std::sync::Arc;use std::time::Duration;use llama_cpp_2::ggml_time_us;use super::args_handler::{Args, Mode};pub struct LLM { pub model: Arc<LlamaModel>, pub backend: LlamaBackend, pub ctx_params: LlamaContextParams, mode: Mode, history: String, max_token: u32, verbose: bool,}new Methodpub fn new(args: Args) -> Self{ // init LLM let backend = LlamaBackend::init() .expect("Could not initialize Llama backend"); // offload all layers to the gpu let model_params = { if args.n_gpu_layers > 0 { LlamaModelParams::default().with_n_gpu_layers(args.n_gpu_layers) } else { LlamaModelParams::default() } }; let mut model_params = pin!(model_params); for (k, v) in &args.key_value_overrides { let k = CString::new(k.as_bytes()).expect(format!("invalid key: {k}").as_str()); model_params.as_mut().append_kv_override(k.as_c_str(), *v); } let model_path = args.model.clone() .get_or_load() .expect("failed to get model from args"); // Load the model and wrap it in an Arc for shared ownership let model = Arc::new(LlamaModel::load_from_file(&backend, model_path, &model_params) .expect("failed to load model" )); // initialize the context let mut ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(args.max_token)) .with_seed(args.seed); if let Some(threads) = args.threads { ctx_params = ctx_params.with_n_threads(threads); } if let Some(threads_batch) = args.threads_batch.or(args.threads) { ctx_params = ctx_params.with_n_threads_batch(threads_batch); } Self { model, backend, ctx_params, mode: args.mode, history: String::new(), max_token: args.max_token, verbose: args.verbose, } }The new method in the LLM struct initializes the core components required for inference, including backend setup, model loading, and context configuration. It takes in an Args instance, which holds the user-provided configuration parameters, and uses these to fine-tune the model and backend settings. Here’s a step-by-step breakdown of each part:Initializing the Backendlet backend = LlamaBackend::init() .expect("Could not initialize Llama backend");Configuring Model Parameterslet model_params = { if args.n_gpu_layers > 0 { LlamaModelParams::default().with_n_gpu_layers(args.n_gpu_layers) } else { LlamaModelParams::default() }};The model_params configuration is determined by whether the user has specified GPU layers with n_gpu_layers. This parameter allows the model to use GPU acceleration for a set number of layers, which can significantly improve inference speed. If no GPU layers are specified (i.e., n_gpu_layers is 0), it defaults to using the CPU.Applying Key-Value Overridesfor (k, v) in &args.key_value_overrides { let k = CString::new(k.as_bytes()).expect(format!("invalid key: {k}").as_str()); model_params.as_mut().append_kv_override(k.as_c_str(), *v);}This section iterates over any key-value overrides specified by the user, allowing them to customize specific model parameters.Loading the Modellet model_path = args.model.clone() .get_or_load() .expect("failed to get model from args");let model = Arc::new(LlamaModel::load_from_file(&backend, model_path, &model_params) .expect("failed to load model"));The model_path is determined by calling get_or_load on the args.model field, which either retrieves the local model path or downloads it from Hugging Face. The model is then loaded into an Arc<LlamaModel> instance, allowing shared ownership of the model.Setting Up the Context Parameterslet mut ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(args.max_token)) .with_seed(args.seed);if let Some(threads) = args.threads { ctx_params = ctx_params.with_n_threads(threads);}if let Some(threads_batch) = args.threads_batch.or(args.threads) { ctx_params = ctx_params.with_n_threads_batch(threads_batch);}Returning the LLM InstanceFinally, the configured fields are assembled into a new LLM instance:Self { model, backend, ctx_params, mode: args.mode, history: String::new(), max_token: args.max_token, verbose: args.verbose,}The history is initialized as an empty string, which is useful for conversational interactions in Chat mode. The max_token and verbose flags are set based on the user input.GenerationThe generate_chat function in this implementation is designed to generate responses in either a chat or text completion mode. It handles various tasks essential for text generation, including formatting inputs, managing tokenization and decoding, and outputting the generated text. Let’s go through each component to see how it works.Formatting the Input Based on ModeThe function starts by formatting the input based on whether the LLM instance is in Chat or Completion mode.let input_to_model = match self.mode { Mode::Chat => { let input_formatted = format!( "<|start_header_id|>user<|end_header_id|>{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>", prompt ); self.history.push_str(input_formatted.as_str()); flag_chat = true; &self.history } Mode::Completion => { &format!("<|begin_of_text|>{}", prompt) }}; Chat Mode: It appends the prompt to a conversation history that includes user and assistant markers. This structured approach keeps track of the conversation context and appends new responses to self.history. Completion Mode: It prepares a straightforward prompt without conversation markers.Tokenizing the InputThe next step tokenizes the formatted prompt, converting it into a list of tokens that the model can process.let tokens_list = ctx.model .str_to_token(&input_to_model, AddBos::Never) .expect(format!("failed to tokenize {}", input_to_model).as_str());Here, AddBos::Never ensures that no beginning-of-sequence token is added, as the prompt format is controlled separately.Validating Token LengthTo prevent the model from running out of memory, the function checks that the token count doesn’t exceed the max_token limit:if tokens_list.len() >= max_token as usize { panic!("the prompt is too long, it has more tokens than max_token ({max_token})")}Preparing the LlamaBatch for DecodingThe function then sets up a LlamaBatch, which is used to manage the tokens being processed.let mut batch = LlamaBatch::new(n_cxt as usize, 1);Each token in tokens_list is added to this batch, with a flag indicating the last token in the sequence for decoding.for (i, token) in tokens_list.into_iter().enumerate() { let is_last = i == last_index; batch.add(token, i as i32, &[0], is_last).unwrap();}Clearing the KV Cache for Completion ModeTo prevent older tokens from influencing new ones in completion mode, the function clears the KV cache:ctx.clear_kv_cache();Decoding LoopThe main loop continues decoding tokens until reaching the max_token limit. During each iteration, it: Samples a New Token: The model generates a list of candidate tokens, from which the most probable token is chosen. let candidates = ctx.candidates(); let candidates_p = LlamaTokenDataArray::from_iter(candidates, false); let new_token_id = ctx.sample_token_greedy(candidates_p); Checks for End of Stream: If the token is an end-of-generation (EOG) marker, the loop breaks. if ctx.model.is_eog_token(new_token_id) { break; } Decodes and Outputs the Token: The selected token ID is converted into a string, appended to the output, and flushed to standard output. let output_bytes = ctx.model.token_to_bytes(new_token_id, Special::Tokenize).unwrap(); let _decode_result = decoder.decode_to_string(&output_bytes, &mut output_string, false); print!("{output_string}"); llm_output.push_str(output_string.as_str()); std::io::stdout().flush().unwrap(); Prepares for Next Iteration: The LlamaBatch is cleared, and the new token is added to continue decoding. Finalizing the OutputAfter decoding, if in chat mode, the output is appended to self.history, preserving context for future exchanges:if flag_chat { self.history.push_str(llm_output_formatted.as_str());}Logging Decoding Speed (Optional)If verbose mode is enabled, the function calculates the speed of token decoding and outputs it.if self.verbose { eprintln!( "\n[decoded {} tokens in {:.2} s, speed {:.2} t/s]", n_decode, duration.as_secs_f32(), n_decode as f32 / duration.as_secs_f32() );}This part is helpful for profiling performance, particularly when optimizing for response time.The main functionThe main function is the entry point for this Rust-based Llama chatbot application. It initializes the model, creates a context for inference, and then enters a loop to handle continuous user input. Let’s break down its components to understand how it operates.Parsing Command-Line Argumentslet args: Args = Args::parse();This line uses the clap library to parse command-line arguments based on the Args struct defined earlier.Initializing the LLM Instancelet mut rllm: LLM = LLM::new(args);Here, we create an instance of LLM using the parsed arguments. The LLM::new function does the heavy lifting of loading the model, setting parameters, and initializing the model backend.Creating the Model Contextlet binding = rllm.model.clone();let mut ctx = binding .new_context(&rllm.backend, rllm.ctx_params.clone()) .expect("failed to create context");The context is necessary for token-based decoding and inference with Llama.The context creation here can feel a bit detached from the LLM struct itself. Since the context setup is specific to a particular input session, it could be made an internal component of LLM, initialized within the struct. But I had lifetime issues!Main Interaction LoopThe function then enters a loop where it continuously takes user input and generates responses until the user exits by submitting an empty input.let mut input = String::new();print!("Assistant: How can I help you today?\n");loop { input.clear(); println!("\nYou: "); std::io::stdout().flush().unwrap(); std::io::stdin().read_line(&mut input).unwrap(); let input = input.trim(); if input.is_empty() { break; } rllm.generate_chat(&mut ctx, &input);} Clearing Input: Before each prompt, input.clear() ensures any leftover data from the previous iteration is removed. Prompting the User: The program prints You: as a prompt to the user to type their input. flush() ensures this prompt displays immediately. Reading and Trimming Input: The user’s input is read from stdin and trimmed to remove extra whitespace. Exiting on Empty Input: If the user submits an empty line, the loop breaks, ending the program. Generating Response: The generate_chat function of LLM is called with the context and user input to generate and display a response from the model.For future refinements, moving context handling inside the LLM struct could streamline the design, making it more self-contained.GitHub RepositoryThe complete implementation is available on GitHub.Reference[1] llama-cpp-rs[2] ChatGPT

Boosting Model Efficiency with Quantization and Pruning in PyTorch

30 minute read

Published: October 01, 2024

Boosting Model Efficiency with Quantization and Pruning in PyTorch The complete implementation is available on GitHub. The complete implementation is available on GitHub.Introduction:In the world of deep learning, deploying models on resource-constrained devices such as mobile phones, IoT devices, or edge computing platforms presents unique challenges. To ensure efficient performance, it’s crucial to reduce model size and enhance inference speed without significantly compromising accuracy. Two popular techniques to achieve this are quantization and pruning.Quantization refers to reducing the precision of the numbers used to represent a model’s weights and activations, typically from 32-bit floating-point numbers to lower precision formats like 8-bit integers. This results in faster computation and lower memory usage. On the other hand, pruning reduces the number of parameters in a model by removing redundant or less important connections, leading to a more compact and faster model.This project showcases a practical implementation of quantization and pruning on a ResNet model using PyTorch. By combining these techniques, we can demonstrate how to create lightweight, efficient models suited for deployment, all while maintaining high accuracy. Whether you are a machine learning practitioner aiming to optimize your models or just exploring model compression techniques, this project serves as an educational guide to help you get started with quantization and pruning in PyTorch.Setup Clone the repository: git clone https://github.com/ramintoosi/resnet-quantization-pruning.git cd resnet-quantization-pruning Install dependencies: pip install torch torchvision torchaudio or conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia Training a Simple Base ModelBefore diving into quantization and pruning, the first step is to train a simple base model without applying any optimization techniques. In this project, we use a standard ResNet architecture to train on the CIFAR-10 dataset. This base model will serve as the foundation for comparing the effects of quantization and pruning later on.The training script, train_simple_model.py, initializes the ResNet model and trains it using typical components such as an Adam optimizer, a cross-entropy loss function, and a learning rate scheduler. By running this script, you will obtain a well-trained model that can later be optimized using various compression techniques.The training process involves the following key steps: Model: A ResNet model with 10 output classes (suitable for CIFAR-10) is instantiated. Device: The model is trained on a GPU if available, or falls back to CPU. Optimizer: The Adam optimizer is used to adjust the model parameters. Learning Rate Scheduler: A ReduceLROnPlateau scheduler is applied to reduce the learning rate when the validation loss stops improving.To train the model from scratch, simply run the following command:python train_model_simple.pyThis will begin the training process and save the best model to weights/simple_best_model.pth. The dataset loading and training loop details are abstracted within the train() function and the load_data() function, which you can refer to in my previous post on weak supervision.After training this base model, we’ll be ready to explore how quantization and pruning can be applied to improve model efficiency.Model ValidationOnce the model is trained, it’s important to evaluate its performance on a validation dataset to measure accuracy, loss, and inference speed. The validation process, implemented in the validate() function, ensures that the model generalizes well to unseen data. Here is the validation script."""This module validates and calculates the accuracy of a model on MNIST validation data."""import timeimport torchfrom tqdm import tqdmfrom data import load_datadef validate(model, device, n_total=2000): """ Validate the model on the validation data. :param n_total: the number of images to validate. :param device: cuda or cpu. :param model: Model to validate. :return: Tuple of accuracy, loss, and average inference time (ms). """ dataloaders = load_data(batch_size=1, num_workers=0) model.eval() model.to(device) correct = 0 total = 0 running_loss = 0.0 start_time = time.time() i_data = 0 with torch.no_grad(): for data in tqdm(dataloaders['val'], total=n_total, desc='Validating model', unit=' image'): images, labels = data[0].to(device), data[1].to(device) outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() running_loss += torch.nn.CrossEntropyLoss()(outputs, labels).item() i_data += 1 if i_data > n_total: break elapsed_time = time.time() - start_time accuracy = correct / total loss = running_loss / total avg_inference_time = elapsed_time / total return accuracy, loss, avg_inference_time * 1000Post-Training Quantization in PyTorchPost-training quantization (PTQ) is a popular technique for reducing the size and improving the efficiency of deep learning models. It allows us to convert a pre-trained floating-point model into a quantized version without requiring retraining. In PyTorch, this is achieved using several quantization strategies, such as dynamic quantization and static quantization. Below is a detailed explanation of the functions used in this module to perform quantization."""This module implements post-training quantization of a PyTorch model."""import copyimport torchfrom torch.ao import quantization as quanimport torch.ao.quantization.quantize_fx as quantize_fxfrom data import load_datafrom tqdm import tqdm# Dynamic Quantization: This method quantizes only the activations during inference, while weights# are quantized beforehand. This means that the quantization overhead occurs during the forward pass, but# since it happens on-the-fly, there's no additional pre-processing step needed.# Therefore, inference time remains unaffected.# This is used for situations where the model execution time is dominated by loading weights# from memory rather than computing the matrix multiplications.def quantize_dynamic(model_fp32: torch.nn.Module, dtype=torch.qint8): """ Quantize a PyTorch model using dynamic quantization. :param model_fp32: model to quantize :param dtype: target dtype for quantized weights :return: quantized model """ # create a quantized model instance model_quantized = torch.ao.quantization.quantize_dynamic( model_fp32, # the original model {torch.nn.Linear}, # a set of layers to dynamically quantize dtype=dtype) # the target dtype for quantized weights return model_quantized# The ModelWrapper class is a custom PyTorch module designed to facilitate the quantization process.# It wraps an existing model and adds QuantStub and DeQuantStub modules to handle the conversion of tensors# between floating point and quantized formats. The forward method specifies where these conversions occur# during the forward pass of the model. This setup is essential for static quantization, where the model needs# to be prepared and calibrated before being converted to a quantized version.class ModelWrapper(torch.nn.Module): def __init__(self, model): super().__init__() # QuantStub converts tensors from floating point to quantized self.quant = torch.ao.quantization.QuantStub() self.model = model # DeQuantStub converts tensors from quantized to floating point self.dequant = torch.ao.quantization.DeQuantStub() def forward(self, x): # manually specify where tensors will be converted from floating # point to quantized in the quantized model x = self.quant(x) x = self.model(x) # manually specify where tensors will be converted from quantized # to floating point in the quantized model x = self.dequant(x) return xdef quantize_static(model: torch.nn.Module): """ Quantize a PyTorch model using static quantization. :param model: model to quantize :return: quantized model """ # wrap the model with the ModelWrapper to include the quant and dequant stubs model_fp32 = ModelWrapper(model) # model must be set to eval mode for static quantization logic to work model_fp32.eval() # attach a global qconfig, which contains information about what kind # of observers to attach. Use 'x86' for server inference and 'qnnpack' # for mobile inference. Other quantization configurations such as selecting # symmetric or asymmetric quantization and MinMax or L2Norm calibration techniques # can be specified here. qconfig = quan.get_default_qconfig('x86') model_fp32.qconfig = qconfig # Prepare the model for static quantization. This inserts observers in # the model that will observe activation tensors during calibration. model_fp32_prepared = torch.ao.quantization.prepare(model_fp32) # calibrate the prepared model to determine quantization parameters for activations # in a real world setting, the calibration would be done with a representative dataset # instead of an empty dataset. dataloader = load_data() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_fp32_prepared.to(device) with torch.no_grad(): for data, _ in tqdm(dataloader['train'], desc='Calibrating model', unit=' batch'): model_fp32_prepared(data.to(device)) model_fp32_prepared.cpu() # Convert the observed model to a quantized model. This does several things: # quantizes the weights, computes and stores the scale and bias value to be # used with each activation tensor, and replaces key operators with quantized # implementations. model_int8 = torch.ao.quantization.convert(model_fp32_prepared) return model_int8def quantize_static_fx(model_fp: torch.nn.Module): """ Quantize a PyTorch model using static quantization with FX graph mode. :param model_fp: model to quantize :return: quantized model """ model_to_quantize = copy.deepcopy(model_fp) qconfig_mapping = quan.get_default_qconfig_mapping("x86") model_to_quantize.eval() # prepare model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, (torch.rand((1, 3, 224, 224)),)) # calibrate dataloader = load_data() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_prepared.to(device) with torch.no_grad(): for data, _ in tqdm(dataloader['train'], desc='Calibrating model FX', unit=' batch'): model_prepared(data.to(device)) model_prepared.cpu() # quantize model_quantized = quantize_fx.convert_fx(model_prepared) return model_quantized1. Dynamic QuantizationDynamic quantization applies quantization only to weights and quantizes activations dynamically during inference. This is useful when the execution time of a model is dominated by memory-bound operations like loading weights.According to the PyTorch website:Dynamic Quantization: This method quantizes only the activations during inference, while weightsare quantized beforehand. This means that the quantization overhead occurs during the forward pass, butsince it happens on-the-fly, there’s no additional pre-processing step needed.Therefore, inference time remains unaffected.This is used for situations where the model execution time is dominated by loading weightsfrom memory rather than computing the matrix multiplications.Function: quantize_dynamic Purpose: This function converts a pre-trained model into a dynamically quantized model by reducing the precision of weights (typically from FP32 to int8). How it works: It takes the floating-point model (model_fp32) and the target data type (dtype, usually torch.qint8). It only quantizes specific layers, such as torch.nn.Linear, where matrix multiplications dominate computation. model_quantized = torch.ao.quantization.quantize_dynamic( model_fp32, # original model {torch.nn.Linear}, # layers to quantize dtype=dtype # target dtype for quantized weights ) 2. Static QuantizationStatic quantization quantizes both weights and activations before inference. It requires calibration to estimate the dynamic range of activations by running representative data through the model.Function: quantize_static Purpose: This function implements static quantization by converting the model into a fully quantized version with both weights and activations quantized. Steps: Model Wrapper: The model is wrapped using ModelWrapper, which introduces QuantStub and DeQuantStub. These stubs manage the conversion between floating-point and quantized tensors during the forward pass. model_fp32 = ModelWrapper(model) Quantization Configuration: The qconfig specifies how the model will be quantized. For example, 'x86' is used for server inference, while qnnpack is suited for mobile devices. qconfig = quan.get_default_qconfig('x86') model_fp32.qconfig = qconfig Prepare for Quantization: The model is prepared for quantization by inserting observers that track the ranges of activations. model_fp32_prepared = torch.ao.quantization.prepare(model_fp32) Calibration: The prepared model is calibrated by passing a subset of the training data through it to estimate the ranges of activations. This step is crucial for determining scaling factors for quantization. for data, _ in tqdm(dataloader['train'], desc='Calibrating model', unit=' batch'): model_fp32_prepared(data.to(device)) Convert to Quantized Model: After calibration, the model is converted to a fully quantized version. This step replaces the floating-point operations with quantized versions. model_int8 = torch.ao.quantization.convert(model_fp32_prepared) 3. FX Graph Mode QuantizationThe FX Graph Mode provides a more flexible approach to quantization by allowing users to modify the model’s computational graph.Function: quantize_static_fx Purpose: This function applies static quantization using PyTorch’s FX Graph Mode API. It allows fine-tuning of the quantization process by manipulating the model’s computational graph. Steps: Model Preparation: The model is deep-copied and prepared for quantization. qconfig_mapping specifies how the layers are quantized, while prepare_fx() converts the model into a graph form suitable for quantization. model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, (torch.rand((1, 3, 224, 224)),)) Calibration: Like static quantization, the model is calibrated by running a sample of training data to estimate the quantization parameters. for data, _ in tqdm(dataloader['train'], desc='Calibrating model FX', unit=' batch'): model_prepared(data.to(device)) Conversion: After calibration, the model is quantized by replacing floating-point operators with their quantized counterparts using convert_fx(). model_quantized = quantize_fx.convert_fx(model_prepared) Summary Dynamic Quantization: Quantizes only weights, with activations quantized on-the-fly during inference. It’s suitable when memory operations dominate. Static Quantization: Quantizes both weights and activations but requires calibration with representative data. It’s ideal for a more compact, efficient model at inference time. FX Graph Mode Quantization: Offers a flexible way to manipulate the model’s computational graph and quantize it with better control.Each quantization method has its own use case depending on the deployment environment and model performance goals.Post-Training Quantization ScriptThis section describes the implementation of post-training quantization in PyTorch, where different quantization techniques are applied to an already trained ResNet model. The primary goal is to compare the accuracy, loss, and inference time of the original and quantized models.The script quantizes the ResNet model using three different PTQ techniques: dynamic quantization, static quantization, and FX static quantization, and then validates the models’ performance."""This module quantizes a PyTorch model using post-training quantization.Let's save the quantized models and compare the results with the original model."""from os.path import isfileimport torchfrom model.resnet import get_modelfrom quantization.post_training import quantize_dynamic, quantize_static, quantize_static_fxfrom validation import validatemodel = get_model(num_classes=10)checkpoint = "weights/original_model.pt"if isfile(checkpoint): model.load_state_dict(torch.load(checkpoint))else: model.load_state_dict(torch.load("weights/simple_best_model.pt")["model_state_dict"]) torch.save(model.state_dict(), "weights/original_model.pt")model_quantized = quantize_dynamic(model, dtype=torch.qint8)checkpoint_quantized = "weights/quantized_dynamic_model.pt"if not isfile(checkpoint_quantized): torch.save(model_quantized.state_dict(), checkpoint_quantized)checkpoint_quantized_static = "weights/quantized_static_model.pt"if isfile(checkpoint_quantized_static): model_quantized_static = torch.jit.load(checkpoint_quantized_static)else: model_quantized_static = quantize_static(model) traced = torch.jit.trace(model_quantized_static, torch.rand((1, 3, 224, 224))) torch.jit.save(traced, checkpoint_quantized_static)checkpoint_quantized_static_fx = "weights/quantized_static_fx_model.pt"if isfile(checkpoint_quantized_static_fx): model_quantized_static_fx = torch.jit.load(checkpoint_quantized_static_fx)else: model_quantized_static_fx = quantize_static_fx(model) traced = torch.jit.trace(model_quantized_static_fx, torch.rand((1, 3, 224, 224))) torch.jit.save(traced, checkpoint_quantized_static_fx)# devicedevice = torch.device("cpu")# validate modelsaccuracy, loss, inference_time = validate(model, device)accuracy_quantized, loss_quantized, inference_time_quantized = validate(model_quantized, device)accuracy_quantized_static, loss_quantized_static, inference_time_quantized_static = ( validate(model_quantized_static, device))accuracy_quantized_static_fx, loss_quantized_static_fx, inference_time_quantized_static_fx = ( validate(model_quantized_static_fx, device))# print the resultsprint(f"Original model accuracy: {accuracy:.4f}, loss: {loss:.4f}, inference time: {inference_time:.2f}ms")print(f"Quantized dynamic model accuracy: {accuracy_quantized:.2f}, loss: {loss_quantized:.2f}, " f"inference time: {inference_time_quantized:.2f}ms")print(f"Quantized static model accuracy: {accuracy_quantized_static:.2f}, loss: {loss_quantized_static:.2f}, " f"inference time: {inference_time_quantized_static:.2f}ms")print(f"Quantized static model with FX accuracy: {accuracy_quantized_static_fx:.2f}, " f"loss: {loss_quantized_static_fx:.2f}, inference time: {inference_time_quantized_static_fx:.2f}ms")Step-by-Step Explanation of the Code Loading the Pre-Trained Model: First, we load a pre-trained ResNet model that has been trained on the CIFAR-10 dataset. model = get_model(num_classes=10) checkpoint = "weights/original_model.pt" if isfile(checkpoint): model.load_state_dict(torch.load(checkpoint)) else: model.load_state_dict(torch.load("weights/simple_best_model.pt")["model_state_dict"]) torch.save(model.state_dict(), "weights/original_model.pt") Dynamic Quantization: model_quantized = quantize_dynamic(model, dtype=torch.qint8) checkpoint_quantized = "weights/quantized_dynamic_model.pt" if not isfile(checkpoint_quantized): torch.save(model_quantized.state_dict(), checkpoint_quantized) Static Quantization: checkpoint_quantized_static = "weights/quantized_static_model.pt" if isfile(checkpoint_quantized_static): model_quantized_static = torch.jit.load(checkpoint_quantized_static) else: model_quantized_static = quantize_static(model) traced = torch.jit.trace(model_quantized_static, torch.rand((1, 3, 224, 224))) torch.jit.save(traced, checkpoint_quantized_static) Static Quantization with FX (Graph Mode): checkpoint_quantized_static_fx = "weights/quantized_static_fx_model.pt" if isfile(checkpoint_quantized_static_fx): model_quantized_static_fx = torch.jit.load(checkpoint_quantized_static_fx) else: model_quantized_static_fx = quantize_static_fx(model) traced = torch.jit.trace(model_quantized_static_fx, torch.rand((1, 3, 224, 224))) torch.jit.save(traced, checkpoint_quantized_static_fx) Model Validation: The original and quantized models are validated on the CIFAR-10 validation dataset using the validate() function. This function calculates the accuracy, loss, and average inference time for each model. accuracy, loss, inference_time = validate(model, device) ... Results Comparison: Finally, the script prints out the results of the original and quantized models, including accuracy, loss, and inference time. print(f"Original model accuracy: {accuracy:.4f}, loss: {loss:.4f}, inference time: {inference_time:.2f}ms") ... Model Pruning FunctionThis section covers the model pruning function, which is used to reduce the number of parameters in a neural network by making the model weights sparse. In this particular implementation, L1 unstructured pruning is applied to both Conv2d and Linear layers of the model.Pruning helps reduce the computational cost and memory footprint of a model, which can be useful for deploying models on resource-constrained environments. Here’s a breakdown of how the pruning function works.import torchimport torch.nn.utils.prune as prunedef make_sparse(model_to_prune, rate=0.5): """ This function prunes the model by making the weights sparse. :param model_to_prune: model to prune :param rate: the percentage of weights to prune """ for name, module in model_to_prune.named_modules(): if isinstance(module, torch.nn.Conv2d): prune.l1_unstructured(module, name='weight', amount=rate) prune.remove(module, 'weight') elif isinstance(module, torch.nn.Linear): prune.l1_unstructured(module, name='weight', amount=rate) prune.remove(module, 'weight')Code Explanation Iterating through the Model’s Layers: The function uses named_modules() to loop through all the layers of the model and check their types. If the layer is a Conv2d or Linear layer, pruning is applied to its weights. L1 Unstructured Pruning: The function applies L1-norm pruning to remove a specified percentage of weights with the smallest absolute values. for name, module in model_to_prune.named_modules(): if isinstance(module, torch.nn.Conv2d): prune.l1_unstructured(module, name='weight', amount=rate) Pruning the Weights: prune.remove() is called afterward to finalize the pruning process and remove the pruning mask from the model. This makes the sparsity permanent, converting zeroed-out weights into actual zeros in the model. prune.remove(module, 'weight') Applying the Same Pruning to Linear Layers: The same pruning process is applied to Linear layers to make them sparse as well. elif isinstance(module, torch.nn.Linear): prune.l1_unstructured(module, name='weight', amount=rate) prune.remove(module, 'weight')Pruning and Quantization Script ExplanationThis script combines pruning and post-training static quantization on a ResNet model to optimize its size and inference performance. Pruning reduces the number of parameters by making the weights sparse, and quantization further compresses the model by reducing the precision of the weights and activations."""This module quantizes a PyTorch model using post-training quantization and pruning."""import copyfrom os.path import isfileimport torchfrom validation import validatefrom model.resnet import get_modelfrom quantization.post_training import quantize_static_fxfrom prune import make_sparsemodel = get_model(num_classes=10)checkpoint = "weights/original_model.pt"model.load_state_dict(torch.load(checkpoint))model_orig = copy.deepcopy(model)checkpoint_quantized_prune = "weights/quantized_prune_model.pt"if isfile(checkpoint_quantized_prune): model_quantized_prune = torch.jit.load(checkpoint_quantized_prune)else: make_sparse(model) model_quantized_prune = quantize_static_fx(model) traced = torch.jit.trace(model_quantized_prune, torch.rand((1, 3, 224, 224))) torch.jit.save(traced, checkpoint_quantized_prune)# validate modelsdevice = torch.device("cpu")accuracy, loss, inference_time = validate(model_orig, device, n_total=100)accuracy_quantized, loss_quantized, inference_time_quantized = validate(model_quantized_prune, device, n_total=100)# print the resultsprint(f"Original model accuracy: {accuracy:.4f}, loss: {loss:.4f}, inference time: {inference_time:.2f}ms")print(f"Quantized static model accuracy: {accuracy_quantized:.2f}, loss: {loss_quantized:.2f}, " f"inference time: {inference_time_quantized:.2f}ms")Here’s a breakdown of each part of the script:Loading the Original Modelmodel = get_model(num_classes=10)checkpoint = "weights/original_model.pt"model.load_state_dict(torch.load(checkpoint))model_orig = copy.deepcopy(model) A ResNet model for 10 classes (e.g., CIFAR-10) is loaded using the get_model() function. The model’s weights are loaded from a saved checkpoint (original_model.pt), ensuring it starts from a pre-trained state. model_orig is a deep copy of the original model, used later for comparison purposes.Pruning and Quantizationcheckpoint_quantized_prune = "weights/quantized_prune_model.pt"if isfile(checkpoint_quantized_prune): model_quantized_prune = torch.jit.load(checkpoint_quantized_prune)else: make_sparse(model) model_quantized_prune = quantize_static_fx(model) traced = torch.jit.trace(model_quantized_prune, torch.rand((1, 3, 224, 224))) torch.jit.save(traced, checkpoint_quantized_prune) The script checks if a quantized and pruned model (quantized_prune_model.pt) already exists. If it exists, the pruned and quantized model is loaded using TorchScript (torch.jit.load()). If not, the following steps occur: Pruning: make_sparse(model) prunes the weights of the model, making the weights sparse by removing a portion of them. Quantization: quantize_static_fx(model) applies static quantization using FX graph mode, further optimizing the model. Saving: The model is traced using TorchScript (torch.jit.trace) for optimization and portability, then saved for future use. Model Validation# validate modelsdevice = torch.device("cpu")accuracy_quantized, loss_quantized, inference_time_quantized = validate(model_quantized_prune, device)Quantization Aware Training (QAT) Script ExplanationThis script demonstrates how to perform Quantization Aware Training (QAT) on a ResNet model using PyTorch. QAT simulates quantization during training, allowing the model to learn to adjust to the quantized weights and activations. This often results in better accuracy when converting the model to a quantized version compared to post-training quantization.Here’s a breakdown of the QAT function and the corresponding QAT training script:QAT Functionfrom torch import nnfrom torch.ao.quantization import get_default_qat_qconfig_mappingimport torch.ao.quantization.quantize_fx as quantize_fximport copydef prepare_model_qat(model_fp: nn.Module, example_inputs): """ Prepare a model for quantization-aware training (QAT). :param model_fp: The floating point model to prepare. :param example_inputs: Example inputs for the model, used during preparation. :return: A model ready for QAT. """ model_to_quantize = copy.deepcopy(model_fp) # Deep copy the model to avoid modifying the original one. # Get the default QAT configuration for the target platform (x86 in this case). qconfig_mapping = get_default_qat_qconfig_mapping("x86") model_to_quantize.train() # Set the model to training mode. # Prepare the model for QAT by adding necessary observers and quantization logic. model_prepared = quantize_fx.prepare_qat_fx(model_to_quantize, qconfig_mapping, example_inputs) return model_prepared Deep Copy of the Model: The function creates a copy of the original floating point model (model_fp) to avoid modifying the original model directly. QAT Config: It retrieves the default QAT configuration for the target platform (x86), specifying how the model’s layers will be quantized during training. Training Mode: The model is set to training mode since QAT requires forward and backward passes to simulate quantization during training. QAT Preparation: The function uses prepare_qat_fx() from the torch.ao.quantization module to prepare the model for quantization-aware training. This adds observers and prepares the model to handle quantized operations during training. QAT Training Script"""this module is used to train the model with QAT"""import osimport torchimport torch.ao.quantization.quantize_fx as quantize_fxfrom data import load_datafrom model.resnet import get_modelfrom train import trainfrom quantization.qat import prepare_model_qatfrom validation import validatedef train_model_qat(resume=True): """ Train a simple model without quantization and pruning. """ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") dataloaders = load_data(batch_size=128, num_workers=0) model = get_model(num_classes=10) model.load_state_dict(torch.load("weights/original_model.pt", weights_only=True)) model_prepared = prepare_model_qat(model, example_inputs = next(iter(dataloaders['train']))[0]) criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.9, patience=5) train(model_prepared, dataloaders, optimizer, criterion, scheduler, device, "qat", 1, resume=resume)if __name__ == '__main__': # train_model_qat(resume=False) checkpoint = 'weights/qat_fx_model.pt' if os.path.isfile(checkpoint): model_quantized = torch.jit.load(checkpoint) else: dataloaders = load_data(batch_size=128, num_workers=0) model = get_model(num_classes=10) model_prepared = prepare_model_qat(model, example_inputs=next(iter(dataloaders['train']))[0]) model_prepared.load_state_dict(torch.load("weights/qat_best_model.pt", weights_only=True)["model_state_dict"]) model_prepared.eval() model_quantized = quantize_fx.convert_fx(model_prepared) traced = torch.jit.trace(model_quantized, torch.rand((1, 3, 224, 224))) torch.jit.save(traced, checkpoint) # validation device = torch.device("cpu") accuracy, loss, inference_time = validate(model_quantized, device) print(f"Quantized model (QAT) accuracy: {accuracy:.2f}, loss: {loss:.2f}, inference time: {inference_time:.2f}ms")Model Training with QATdef train_model_qat(resume=True): """ Train a ResNet model with Quantization Aware Training (QAT). """ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") dataloaders = load_data(batch_size=128, num_workers=0) model = get_model(num_classes=10) model.load_state_dict(torch.load("weights/original_model.pt", weights_only=True)) model_prepared = prepare_model_qat(model, example_inputs=next(iter(dataloaders['train']))[0]) criterion = torch.nn.CrossEntropyLoss() # Loss function optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # Adam optimizer scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.9, patience=5) # Learning rate scheduler train(model_prepared, dataloaders, optimizer, criterion, scheduler, device, "qat", 1, resume=resume) The model is loaded with the original pre-trained weights. The model is prepared for QAT using the prepare_model_qat() function with example inputs from the training data. The script sets up the loss function, optimizer, and learning rate scheduler to adjust the learning rate when the model’s performance plateaus.Loading and Quantizing the Modelif __name__ == '__main__': checkpoint = 'weights/qat_fx_model.pt' if os.path.isfile(checkpoint): model_quantized = torch.jit.load(checkpoint) # Load the pre-quantized model else: dataloaders = load_data(batch_size=128, num_workers=0) model = get_model(num_classes=10) model_prepared = prepare_model_qat(model, example_inputs=next(iter(dataloaders['train']))[0]) model_prepared.load_state_dict(torch.load("weights/qat_best_model.pt", weights_only=True)["model_state_dict"]) model_prepared.eval() # Set the model to evaluation mode before quantizing. model_quantized = quantize_fx.convert_fx(model_prepared) # Convert the model to quantized format. traced = torch.jit.trace(model_quantized, torch.rand((1, 3, 224, 224))) # Trace the model with TorchScript. torch.jit.save(traced, checkpoint) # Save the quantized model. If a pre-quantized model checkpoint exists, it is loaded directly. Otherwise, the script prepares the model for QAT and loads the trained QAT model’s weights. It then converts the QAT-trained model to a fully quantized format using quantize_fx.convert_fx(). The quantized model is traced with TorchScript and saved to be loaded later.Model Validation# validationdevice = torch.device("cpu")accuracy, loss, inference_time = validate(model_quantized, device)print(f"Quantized model (QAT) accuracy: {accuracy:.2f}, loss: {loss:.2f}, inference time: {inference_time:.2f}ms")Pruning and QAT"""this module is used to train the model with QAT"""import osimport torchimport torch.ao.quantization.quantize_fx as quantize_fxfrom data import load_datafrom model.resnet import get_modelfrom train import trainfrom quantization.qat import prepare_model_qatfrom validation import validatedef train_model_qat(resume=True): """ Train a simple model without quantization and pruning. """ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") dataloaders = load_data(batch_size=128, num_workers=0) model = get_model(num_classes=10) model.load_state_dict(torch.load("weights/original_model.pt", weights_only=True)) model_prepared = prepare_model_qat(model, example_inputs = next(iter(dataloaders['train']))[0]) criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.9, patience=5) train(model_prepared, dataloaders, optimizer, criterion, scheduler, device, "qat", 1, resume=resume)if __name__ == '__main__': # train_model_qat(resume=False) checkpoint = 'weights/qat_fx_model.pt' if os.path.isfile(checkpoint): model_quantized = torch.jit.load(checkpoint) else: dataloaders = load_data(batch_size=128, num_workers=0) model = get_model(num_classes=10) model_prepared = prepare_model_qat(model, example_inputs=next(iter(dataloaders['train']))[0]) model_prepared.load_state_dict(torch.load("weights/qat_best_model.pt", weights_only=True)["model_state_dict"]) model_prepared.eval() model_quantized = quantize_fx.convert_fx(model_prepared) traced = torch.jit.trace(model_quantized, torch.rand((1, 3, 224, 224))) torch.jit.save(traced, checkpoint) # validation device = torch.device("cpu") accuracy, loss, inference_time = validate(model_quantized, device) print(f"Quantized model (QAT) accuracy: {accuracy:.2f}, loss: {loss:.2f}, inference time: {inference_time:.2f}ms")This script performs both pruning and Quantization Aware Training (QAT) on a ResNet model using PyTorch. By combining pruning, which reduces the number of weights, and QAT, which trains the model to adjust for quantization, we achieve both model compression and faster inference while maintaining high accuracy.Here’s a breakdown of the Pruning and QAT function and the corresponding training script:Code BreakdownModel Training with Pruning and QATdef train_model_qat_prune(resume=True): """ Train a ResNet model with Quantization Aware Training (QAT) and pruning. """ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") dataloaders = load_data(batch_size=128, num_workers=0) model = get_model(num_classes=10) model.load_state_dict(torch.load("weights/original_model.pt", weights_only=True)) # Apply pruning to make the model sparse make_sparse(model) # Prepare the pruned model for Quantization Aware Training (QAT) model_prepared = prepare_model_qat(model, example_inputs=next(iter(dataloaders['train']))[0]) criterion = torch.nn.CrossEntropyLoss() # Define loss function optimizer = torch.optim.Adam(model.parameters(), lr=0.0005) # Adam optimizer with a smaller learning rate scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.9, patience=5) # Learning rate scheduler # Train the pruned and QAT-prepared model train(model_prepared, dataloaders, optimizer, criterion, scheduler, device, "qat_prune", 20, resume=resume) Pruning: The make_sparse(model) function prunes the ResNet model by setting a percentage of the model’s weights to zero, based on the pruning rate. This function is applied to the model before QAT. QAT Preparation: The pruned model is passed through the prepare_model_qat() function, which prepares it for Quantization Aware Training by inserting quantization-aware operations into the model. This prepares the model to handle quantized weights during training. Training Setup: The script defines the training setup using CrossEntropyLoss, the Adam optimizer, and a learning rate scheduler. The learning rate is reduced when the model’s performance plateaus. Training: The training loop is executed for 20 epochs ("qat_prune", 20), adjusting the weights of the pruned and quantization-aware model. Model Loading and Quantizationif __name__ == '__main__': train_model_qat_prune(resume=False) checkpoint = 'weights/qat_prune_model.pt' if os.path.isfile(checkpoint): model_quantized = torch.jit.load(checkpoint) # Load the pre-quantized model if it exists else: dataloaders = load_data(batch_size=128, num_workers=0) model = get_model(num_classes=10) # Apply pruning make_sparse(model) # Prepare the pruned model for QAT model_prepared = prepare_model_qat(model, example_inputs=next(iter(dataloaders['train']))[0]) model_prepared.load_state_dict(torch.load("weights/qat_prune_best_model.pt", weights_only=True)["model_state_dict"]) # Convert the QAT-prepared model to a quantized format model_prepared.eval() model_quantized = quantize_fx.convert_fx(model_prepared) # Trace and save the quantized model traced = torch.jit.trace(model_quantized, torch.rand((1, 3, 224, 224))) torch.jit.save(traced, checkpoint) Loading Pre-Trained Model: Pruning and QAT: If no checkpoint exists, the model is pruned, and then it is prepared for Quantization Aware Training with prepare_model_qat(). After loading the pre-trained weights (weights/qat_prune_best_model.pt), the model is converted to its fully quantized form using quantize_fx.convert_fx(). Model Tracing and Saving:Model Validation # Validation step device = torch.device("cpu") accuracy, loss, inference_time = validate(model_quantized, device) print(f"Quantized model (QAT) accuracy: {accuracy:.2f}, loss: {loss:.2f}, inference time: {inference_time:.2f}ms")Results SummaryThese results show the comparison of various model optimization techniques applied to a ResNet model trained on CIFAR-10. The experiments were conducted using an NVIDIA RTX 2080 Ti GPU, and the results focus on three metrics: Model Type Accuracy Loss Inference Time Original model 0.95 0.27 54.28ms PTQ dynamic model 0.96 0.27 53.95ms PTQ static model 0.95 0.28 22.96ms PTQ static model with FX 0.95 0.28 21.37ms Pruned 50% + PTQ static 0.93 0.19 20.02ms QAT 0.95 0.27 19.87ms Pruned 50% + QAT 0.95 0.25 20.61ms Key Observations Accuracy: The original model achieves 95% accuracy, and both PTQ (Post-Training Quantization) and QAT (Quantization Aware Training) retain similar levels of accuracy (95%-96%), demonstrating minimal accuracy loss. Pruned 50% + PTQ static slightly reduces accuracy to 93%, but applying QAT along with pruning helps recover accuracy to 95%. Inference Time: PTQ static models achieve a significant reduction in inference time, cutting it by over 50% (from 54.28ms to ~21ms). Pruning and QAT further reduce the inference time, with QAT alone achieving the fastest inference at 19.87ms. Pruned 50% + PTQ/QAT models maintain a low inference time of around 20ms, which is a substantial improvement over the original model. Loss: The original model, PTQ, and QAT maintain similar loss values (~0.27-0.28). Pruned models show a lower loss (0.19-0.25), suggesting that pruning can help reduce overfitting and potentially improve generalization. Conclusion This experiment demonstrates the effectiveness of post-training quantization (PTQ), quantization-aware training (QAT), and pruning techniques in optimizing the performance of a deep learning model. By applying these techniques to a ResNet model trained on CIFAR-10, we observed substantial improvements in inference time and model efficiency while maintaining accuracy levels close to the original, unoptimized model.Both static and dynamic quantization techniques, particularly static quantization with FX, significantly reduced inference time without sacrificing accuracy. Similarly, pruning, especially when combined with PTQ or QAT, further enhanced model efficiency while maintaining strong performance metrics.QAT emerged as a standout technique for producing highly efficient models, yielding minimal accuracy loss and achieving the fastest inference time of 19.87ms. The pruned models also performed well, showing that combining QAT or PTQ with pruning is an effective strategy for improving both model speed and size.Takeaways Quantization is highly effective in optimizing model performance: QAT achieves the best balance between speed and accuracy: Pruning enhances the benefits of quantization: Pruning alone can reduce overfitting, as seen in the lower loss values. Static quantization offers superior speedup compared to dynamic quantization: Pruned models show a slight drop in accuracy but perform well overall:Final Thought:For real-world applications requiring model optimization, QAT combined with pruning offers the best balance of speed, accuracy, and efficiency. When ease of implementation is a concern, PTQ static is a strong alternative that still provides significant benefits in terms of inference time.GitHub RepositoryThe complete implementation is available on GitHub.Reference[1] PyTorch Quantization[2] PyTorch Pruning[3] ChatGPT

Let’s train a CNN in Rust with libtorch

18 minute read

Published: June 29, 2024

Let's train a CNN in Rust with libtorch Introduction: Introduction:As a newcomer to Rust, I recently took on the challenge of training a CNN to classify images of the sea versus the jungle using the tch-rust crate, which employs Libtorch. Rust’s promise of memory safety and performance intrigued me, and I was eager to explore how it could handle deep learning tasks. In this post, I’ll walk you through my journey of setting up and training a simple CNN model. Whether you’re a seasoned Rustacean or just starting, I hope this post will inspire you to experiment with deep learning in Rust.Note: I’m not an experienced Rust developer.DatasetFor the dataset, I used a collection I prepared in a previous blog post. You can use any dataset that is structured as follows:data/ ├── train/ │ ├── class1/ │ │ ├── img1.jpg │ │ ├── img2.jpg │ ├── class2/ │ │ ├── img1.jpg │ │ ├── img2.jpg ├── val/ ├── class1/ │ ├── img1.jpg │ ├── img2.jpg ├── class2/ ├── img1.jpg ├── img2.jpgThis structure makes it easy to load images for training and validation, with separate folders for each class.libtorchWe use the tch-rust crate, which provides “Rust bindings for the C++ API of PyTorch.” This crate requires Libtorch. You can download pre-built Libtorch files from PyTorch or build them manually. Then you need to set two env variables.export LIBTORCH=[path_to_libroch]export LD_LIBRARY_PATH=[path_to_libtorch]:$LD_LIBRARY_PATHpath_to_libtorch is the Libtorch folder. If you downloaded the pre-built files, use that path. If you built it manually, it might look something like /path/to/pytorch/build/lib.linux-x86_64-cpython-310/torch/.Build libtorch manuallyInstall typing-extensions and pyyaml. Then,git clone -b v2.3.0 --recurse-submodule https://github.com/pytorch/pytorch.git --depth 1cd pytorchUSE_CUDA=ON BUILD_SHARED_LIBS=ON python setup.py buildChoose pytorch version based on the tch-rust crate.Code structureHere’s a brief overview of the project file structure:├── src│ ├── train│ │ └── utils.rs # Utility functions for training│ ├── data.rs # Data loading and preprocessing│ ├── inference.rs # Functions for model inference│ ├── main.rs # Main program entry point│ ├── model.rs # Definition of the CNN model│ └── train.rs # Training loop and logic main.rs: The main entry point of the application, where execution begins. data.rs: Handles data loading and preprocessing, organizing datasets for training and validation. model.rs: Defines the architecture of the CNN model. train.rs: Implements the training loop and logic, coordinating the model training process. train/utils.rs: Contains helper functions used during training. inference.rs: Includes functions for running inference on new images.Data loaderThe data loader module is designed to handle the image dataset, similar to ImageFolderDataset in PyTorch. It assumes that images are divided into training and validation sets, with the images of each class in its own folder.use kdam::tqdm;use rand::seq::SliceRandom;use rand::thread_rng;use std::collections::HashMap;use std::{fs::read_dir, path::Path};use tch::{vision, Tensor};pub struct Dataset { root: String, image_path: Vec<(i64, String)>, class_to_idx: HashMap<String, i64>, total_size: usize,}impl Dataset { /// This function walks through the root folder and gathers images and creates a Dataset pub fn new<T: AsRef<Path>>(root: T) -> Dataset { let root = root.as_ref(); let mut image_path: Vec<(i64, String)> = Vec::new(); let mut class_to_idx: HashMap<String, i64> = HashMap::new(); Self::get_images_and_classes( &root, &mut image_path, &mut class_to_idx, ); Dataset { root: root.to_str().unwrap().to_string(), total_size: image_path.len(), image_path, class_to_idx, } } /// In the input folder finds the classes and images fn get_images_and_classes( dir: &Path, image_path: &mut Vec<(i64, String)>, class_to_idx: &mut HashMap<String, i64>, ) { for (class_id, root_class) in read_dir(&dir).unwrap().enumerate() { let root_class = root_class.unwrap().path().clone(); if root_class.is_dir() { Self::get_images_in_folder(&root_class, image_path, class_id as i64); let class_name_str = root_class .file_name() .unwrap() .to_str() .unwrap() .to_string(); class_to_idx.insert(class_name_str.clone(), class_id as i64); } } } /// find images with specific extensions "jpg", "png", "jpeg" fn get_images_in_folder( dir: &Path, image_path: &mut Vec<(i64, String)>, class_idx: i64, ) { let valid_ext = vec!["jpg", "png", "jpeg"]; for file_path in tqdm!(read_dir(&dir).unwrap()) { let file_path = &file_path.unwrap().path().clone(); if file_path.is_file() & valid_ext.contains( &file_path .extension() .unwrap() .to_str() .unwrap() .to_lowercase() .as_str(), ) { image_path.push((class_idx, file_path.to_str().unwrap().to_string())); } } } /// A simple print function for our Dataset pub fn print(&self) { println!("DATASET ({})", self.root); println!("Classes: {:?}", self.class_to_idx); println!("Size: {}", self.total_size); println!("sample of data\n{:?}", &self.image_path[1..3]); } /// load the image into a tensor and return (image, label) fn get_item(&self, idx: usize) -> (Tensor, i64) { let image =vision::imagenet::load_image_and_resize224(&self.image_path[idx].1).unwrap(); (image, self.image_path[idx].0.clone()) }}/// A struct for our data loaderpub struct DataLoader { dataset: Dataset, batch_size: i64, batch_index: i64, shuffle: bool,}impl DataLoader { pub fn new(dataset: Dataset, batch_size: i64, shuffle: bool) -> DataLoader { // let mut rng = thread_rng(); // dataset.ImagePath.shuffle(rng); DataLoader { dataset, batch_size, batch_index: 0, shuffle, } } fn shuffle_dataset(&mut self) { let mut rng = thread_rng(); self.dataset.image_path.shuffle(&mut rng) } /// total number of images in the dataset pub fn len(&self) -> usize { self.dataset.total_size } /// number of batches based on the dataset size and batch size pub fn len_batch(&self) -> usize { (self.dataset.total_size / self.batch_size as usize) + 1 }}/// implement iterator for our Dataloader to get batches of images and labelsimpl Iterator for DataLoader { type Item = (Tensor, Tensor); fn next(&mut self) -> Option<Self::Item> { let start = (self.batch_index * self.batch_size) as usize; let mut end = ((self.batch_index + 1) * self.batch_size) as usize; if start >= self.dataset.total_size { self.batch_index = 0; return None; } if end > self.dataset.total_size { end = self.dataset.total_size; } if (self.batch_index == 0) & self.shuffle { self.shuffle_dataset(); } let mut images: Vec<Tensor> = vec![]; // for preload change this to Vec<&Tensor> let mut labels: Vec<Tensor> = vec![]; for i in start..end { let (image_t, label) = self.dataset.get_item(i); images.push(image_t); labels.push(Tensor::from(label)) } self.batch_index += 1; Some(( Tensor::f_stack(&images, 0).unwrap(), Tensor::f_stack(&labels, 0).unwrap(), )) }}Dataset: Manages the dataset structure, storing image paths and class indices. It reads the directory to gather images and organizes them by class.pub struct Dataset { root: String, image_path: Vec<(i64, String)>, class_to_idx: HashMap<String, i64>, total_size: usize,}Key Functions new: Initializes the dataset by reading image paths and classes from the root directory. get_images_and_classes: Recursively finds class directories and images within them. get_images_in_folder: Filters images based on valid extensions (jpg, png, jpeg). print: Outputs dataset information, including classes and sample data. get_item: Loads an image as a tensor and returns it along with its label.DataLoader: Handles batching of the dataset for training. It can shuffle data and iterates through batches of images and labels.pub struct DataLoader { dataset: Dataset, batch_size: i64, batch_index: i64, shuffle: bool,}Key Functions new: Initializes the data loader with the specified batch size and shuffle option. shuffle_dataset: Randomly shuffles the dataset. len: Returns the total number of images. len_batch: Calculates the total number of batches.Iterator ImplementationThe DataLoader implements the Iterator trait to provide batches of images and labels for training: next: Retrieves the next batch, stacking images and labels into tensors. If all batches are processed, it resets and optionally shuffles the dataset for the next epoch.This module organizes and processes image data, facilitating smooth training with Rust and Libtorch.Model DefinitionHere’s the definition of our CNN model using the tch crate. The model consists of several convolutional layers followed by fully connected layers.use tch::{nn, nn::Module};pub fn net(vs: &nn::Path, n_class: i64) -> impl Module { nn::seq() .add(nn::conv2d(vs, 3, 16, 16, Default::default())) .add_fn(|xs| xs.max_pool2d_default(4)) .add_fn(|xs| xs.relu()) .add(nn::conv2d(vs, 16, 64, 4, Default::default())) .add_fn(|xs| xs.max_pool2d_default(2)) .add_fn(|xs| xs.relu()) .add(nn::conv2d(vs, 64, 128, 4, Default::default())) .add_fn(|xs| xs.relu()) .add_fn(|xs| xs.flat_view()) .add(nn::linear(vs, 56448, 1024, Default::default())) .add_fn(|xs| xs.relu()) .add(nn::linear(vs, 1024, n_class, Default::default()))}Model Architecture Convolutional Layers: Three convolutional layers with increasing channel sizes (16, 64, 128) to capture spatial features. Each is followed by a ReLU activation and pooling layer to reduce dimensionality and introduce non-linearity. Fully Connected Layers: Two fully connected layers that transform the features into the desired number of classes. The first linear layer reduces the feature map size from 56448 to 1024. The final layer outputs predictions for each class. Learning Rate SchedulerThe scheduler dynamically adjusts the learning rate during training based on model performance. This helps optimize convergence.use num_traits::Float;use tch::nn::Optimizer;pub struct Scheduler<'a> { pub opt: &'a mut Optimizer, patience: i64, factor: f64, lr: f64, step: i64, last_val: f64,}impl Scheduler<'_> { pub fn new(opt: &mut Optimizer, mut patience: i64, lr: f64, mut factor: f64) -> Scheduler { if patience < 0 { patience = 5; } if factor < 0.0 { factor = 0.95; } Scheduler { opt, patience, factor, lr, step: 0, last_val: f64::infinity(), } } /// Adjusts learning rate based on validation performance. pub fn step(&mut self, value: f64) { if value < self.last_val { self.last_val = value; self.step = 0; } else { self.step += 1; if self.step == self.patience { self.step = 0; self.lr *= self.factor; self.opt.set_lr(self.lr); } } } pub fn get_lr(&self) -> f64 { self.lr }}Scheduler Functionality Initialization: Takes an optimizer, patience, initial learning rate, and a decay factor. Default values are set if negative inputs are given. step Method: Monitors validation loss to decide when to reduce the learning rate. If the validation loss improves, it resets the step counter; otherwise, it increments the counter. After reaching the patience threshold, the learning rate is decreased by the factor. get_lr Method: Returns the current learning rate.This scheduler helps fine-tune the training process by adapting the learning rate, ensuring efficient convergence.Model Training Step-by-Step Set Up Directories and Device: Check if the save_dir exists; create it if not. Determine if CUDA is available and select the appropriate device for training. if !Path::new(save_dir).is_dir() { create_dir(save_dir).unwrap(); } let device = Device::cuda_if_available(); println!("The device is {:?}", device); Initialize Model and Optimizer: Create the model using model::net with two output classes. Set up an Adam optimizer with a learning rate of 1e-3 and a weight decay of 1e-4. let vs = nn::VarStore::new(device); let net = model::net(&vs.root(), 2); let lr = 1e-3; let mut opt = nn::Adam::default().build(&vs, lr).unwrap(); opt.set_weight_decay(1e-4); Learning Rate Scheduler: Initialize a scheduler to adjust the learning rate dynamically based on validation performance, with patience set to 5 epochs and a decay factor of 0.5. let mut scheduler = utils::Scheduler::new( &mut opt, 5, lr, 0.5 ); Progress Bars: Set up progress bars for tracking training and validation progress, as well as the overall epoch. let total_batch_train = dataloader_train.len_batch(); let total_batch_val = dataloader_val.len_batch(); let mut pbar = tqdm!( total = total_batch_train, position = 1, desc = format!("{:<8}", "Train"), force_refresh = true, ncols = 100 ); let mut pbar2 = tqdm!( total = total_batch_val, position = 2, desc = format!("{:<8}", "Val"), force_refresh = true, ncols = 100 ); let n_epochs = 30; let mut pbar_e = tqdm!( total = n_epochs, position = 0, desc = format!("{:<8}", "Epoch"), ncols = 100, force_refresh=true ); Training Loop: Run for a fixed number of epochs (30). For each epoch: Training Phase: Iterate over the training data loader. For each batch: Zero the gradients. Forward pass: compute predictions. Calculate accuracy and cross-entropy loss. Backward pass: compute gradients and update weights. Update the progress bar with the current loss and accuracy. Validation Phase: Iterate over the validation data loader. For each batch: Forward pass: compute predictions. Calculate accuracy and cross-entropy loss. Update the progress bar with the current loss and accuracy. Adjust Learning Rate: After each epoch, the scheduler checks the validation loss to decide if the learning rate should be decreased. Save Best Model: If the current epoch’s validation loss is the lowest seen so far, save the model. let mut best_acc = 0.0;let mut best_loss: f64 = f64::infinity();println!("\n\n Start Training \n\n");for e in 1..n_epochs { pbar_e.set_postfix( format!("lr = {:<.7}", scheduler.get_lr()) ); let _ = pbar_e.update_to(e); let mut epoch_acc_train = 0.0; let mut epoch_loss_train = 0.0; let mut running_samples = 0; for (i, (images, labels)) in (&mut dataloader_train).enumerate() { scheduler.opt.zero_grad(); let out = net .forward(&images.to_device(device)) .to_device(Device::Cpu); let acc = out.accuracy_for_logits(&labels); let loss = out.cross_entropy_for_logits(&labels); epoch_acc_train += f64::try_from(acc).unwrap() * (out.size()[0] as f64); epoch_loss_train += f64::try_from(&loss).unwrap() * (out.size()[0] as f64); scheduler.opt.backward_step(&loss); running_samples += out.size()[0]; pbar.set_postfix(format!( "loss={:<7.4} - accuracy={:<7.4}", epoch_loss_train / (running_samples as f64), epoch_acc_train / (running_samples as f64) * 100.0 )); let _ = pbar.update_to(i + 1); } let mut epoch_acc_val = 0.0; let mut epoch_loss_val = 0.0; running_samples = 0; for (i, (images, labels)) in (&mut dataloader_val).enumerate() { let out = net .forward(&images.to_device(device)) .to_device(Device::Cpu); let loss = out.cross_entropy_for_logits(&labels); let acc = out.accuracy_for_logits(&labels); epoch_acc_val += f64::try_from(acc).unwrap() * (out.size()[0] as f64); epoch_loss_val += f64::try_from(&loss).unwrap() * (out.size()[0] as f64); running_samples += out.size()[0]; pbar2.set_postfix(format!( "loss={:<7.4} - accuracy={:<7.4}", epoch_loss_val / (running_samples as f64), epoch_acc_val / (running_samples as f64) * 100.0 )); let _ = pbar2.update_to(i + 1); } epoch_acc_val /= dataloader_val.len() as f64; epoch_loss_val /= dataloader_val.len() as f64; scheduler.step(epoch_loss_val); if epoch_loss_val < best_loss { best_loss = epoch_loss_val; best_acc = epoch_acc_val; vs.save(Path::new(save_dir).join("best_model.ot")).unwrap() }}Key Points Device Management: Utilizes GPU if available for faster computations. Model and Optimizer: Defined using the tch crate, with appropriate configurations. Dynamic Learning Rate: Adjusts based on validation performance to optimize training. Progress Tracking: Uses kdam crate to visually track training and validation progress. Performance Metrics: Tracks accuracy and loss during both training and validation phases.Complete training codemod utils;use crate::data::DataLoader;use crate::model;use kdam::{tqdm, BarExt};use num_traits::float::Float;use tch::{ nn, nn::{Module, OptimizerConfig}, Device,};use std::path::Path;use std::fs::create_dir;/// This function trains the model with train and val data loaderspub fn train_model( mut dataloader_train: DataLoader, mut dataloader_val: DataLoader, save_dir: &str,) { if !Path::new(save_dir).is_dir() { create_dir(save_dir).unwrap(); } let device = Device::cuda_if_available(); println!("The device is {:?}", device); let vs = nn::VarStore::new(device); let net = model::net(&vs.root(), 2); let lr = 1e-3; let mut opt = nn::Adam::default().build(&vs, lr).unwrap(); opt.set_weight_decay(1e-4); let mut scheduler = utils::Scheduler::new( &mut opt, 5, lr, 0.5 ); let total_batch_train = dataloader_train.len_batch(); let total_batch_val = dataloader_val.len_batch(); let mut pbar = tqdm!( total = total_batch_train, position = 1, desc = format!("{:<8}", "Train"), force_refresh = true, ncols = 100 ); let mut pbar2 = tqdm!( total = total_batch_val, position = 2, desc = format!("{:<8}", "Val"), force_refresh = true, ncols = 100 ); let n_epochs = 30; let mut pbar_e = tqdm!( total = n_epochs, position = 0, desc = format!("{:<8}", "Epoch"), ncols = 100, force_refresh=true ); let mut best_acc = 0.0; let mut best_loss: f64 = f64::infinity(); println!("\n\n Start Training \n\n"); for e in 1..n_epochs { pbar_e.set_postfix( format!("lr = {:<.7}", scheduler.get_lr()) ); let _ = pbar_e.update_to(e); let mut epoch_acc_train = 0.0; let mut epoch_loss_train = 0.0; let mut running_samples = 0; for (i, (images, labels)) in (&mut dataloader_train).enumerate() { scheduler.opt.zero_grad(); let out = net .forward(&images.to_device(device)) .to_device(Device::Cpu); let acc = out.accuracy_for_logits(&labels); let loss = out.cross_entropy_for_logits(&labels); epoch_acc_train += f64::try_from(acc).unwrap() * (out.size()[0] as f64); epoch_loss_train += f64::try_from(&loss).unwrap() * (out.size()[0] as f64); scheduler.opt.backward_step(&loss); running_samples += out.size()[0]; pbar.set_postfix(format!( "loss={:<7.4} - accuracy={:<7.4}", epoch_loss_train / (running_samples as f64), epoch_acc_train / (running_samples as f64) * 100.0 )); let _ = pbar.update_to(i + 1); } let mut epoch_acc_val = 0.0; let mut epoch_loss_val = 0.0; running_samples = 0; for (i, (images, labels)) in (&mut dataloader_val).enumerate() { let out = net .forward(&images.to_device(device)) .to_device(Device::Cpu); let loss = out.cross_entropy_for_logits(&labels); let acc = out.accuracy_for_logits(&labels); epoch_acc_val += f64::try_from(acc).unwrap() * (out.size()[0] as f64); epoch_loss_val += f64::try_from(&loss).unwrap() * (out.size()[0] as f64); running_samples += out.size()[0]; pbar2.set_postfix(format!( "loss={:<7.4} - accuracy={:<7.4}", epoch_loss_val / (running_samples as f64), epoch_acc_val / (running_samples as f64) * 100.0 )); let _ = pbar2.update_to(i + 1); } epoch_acc_val /= dataloader_val.len() as f64; epoch_loss_val /= dataloader_val.len() as f64; scheduler.step(epoch_loss_val); if epoch_loss_val < best_loss { best_loss = epoch_loss_val; best_acc = epoch_acc_val; vs.save(Path::new(save_dir).join("best_model.ot")).unwrap() } } println!("\n\n\n"); println!( "Best validation loss = {best_loss:.4}, accuracy={:.4}", best_acc * 100.0 );}InferenceThe inference function predicts the class of an input image using a trained CNN model. It first checks for CUDA availability and sets the device. The model is initialized and loaded with pre-trained weights. The input image is preprocessed, resized to 224x224 pixels, and passed through the model. The function then identifies the predicted class by selecting the index of the maximum output logit, returning this as the predicted class label. This process enables quick classification of new images.use tch::{Device, nn, vision, nn::Module};use crate::model;pub fn inference(image_path: &str) -> i64{ let device = Device::cuda_if_available(); let mut vs = nn::VarStore::new(device); let net = model::net(&vs.root(), 2); vs.load("weights/best_model.ot").unwrap(); let image = vision::imagenet::load_image_and_resize224(image_path).unwrap().unsqueeze(0); let out = net.forward(&image.to_device(device)); let prediction = out.argmax(1, false); i64::try_from(prediction).unwrap()}Main function: A Complete WorkflowThe main function orchestrates the training and inference processes of our CNN. The datasets for training and validation are loaded from specified directories, and their details are printed, including class mappings and dataset sizes. Data loaders are then initialized for batch processing during training and validation phases. The CNN model is trained using the training dataset, with validation performed using the separate validation dataset to monitor performance. After training, the best-performing model weights are saved. Inference is demonstrated by predicting the class of a sample image from the validation set using the trained model, and the predicted class label is printed as output.mod data;mod model;mod train;mod inference;fn main() { std::env::set_var("CUDA_LAUNCH_BLOCKING", "1"); std::env::set_var("TORCH_SHOW_WARNINGS", "0"); println!("Loading Dataset / train"); let dataset_train = data::Dataset::new("data/sea_vs_jungle/train"); println!("Loading Dataset / val"); let dataset_val = data::Dataset::new("data/sea_vs_jungle/val"); dataset_train.print(); dataset_val.print(); let dataloader_train = data::DataLoader::new(dataset_train, 32, true); let dataloader_val = data::DataLoader::new(dataset_val, 32, false); println!("{}", dataloader_val.len_batch()); train::train_model(dataloader_train, dataloader_val, "weights"); let prediction = inference::inference("data/sea_vs_jungle/val/sea/001c31c29de8a9cd.jpg"); println!("Prediction is {prediction}")}ConclusionAt the end of this project, while Rust proved to be a powerful tool for deep learning tasks with LibTorch bindings using the tch-rs crate, I found that it might be more straightforward to train a model in Python and then utilize it in Rust. In the future, I plan to explore this approach further and will write a blog post on how to train a model in Python, extract it using TorchScript, and seamlessly integrate it into Rust applications.Code RepositoryThe complete Rust implementation is available on GitHub.Reference[1] Rust Book[2] tch-rust[3] ChatGPT

Implementing K-Means in Rust: A 7x Speed Boost Over Python

10 minute read

Published: April 06, 2024

Implementing K-Means in Rust: A 7x Speed Boost Over Python Introduction:As a machine learning enthusiast, I’ve always been fascinated by the speed and efficiency of algorithms. Recently, I decided to learn about Rust, a language known for its performance and safety guarantees, to implement the K-means clustering algorithm. Despite not being an experienced Rust developer, I was amazed by the results: my Rust implementation outperformed the popular scikit-learn package in Python by a factor of 7. In this blog post, I’ll share my journey of implementing K-means in Rust. Let’s dive in! Introduction:As a machine learning enthusiast, I’ve always been fascinated by the speed and efficiency of algorithms. Recently, I decided to learn about Rust, a language known for its performance and safety guarantees, to implement the K-means clustering algorithm. Despite not being an experienced Rust developer, I was amazed by the results: my Rust implementation outperformed the popular scikit-learn package in Python by a factor of 7. In this blog post, I’ll share my journey of implementing K-means in Rust. Let’s dive in!DataThis Python code generates 1 million samples of Gaussian noise clusters with different means in 10 dimensions and saves them to a CSV file named “data.csv”. There are four clusters, and each cluster has different means, and the clusters are separated by a distance of 2 along each dimension. Here’s the code:import numpy as npn_cluster = 4number_of_samples = 1_000_000dim = 10n_sample_per_class = int(number_of_samples / n_cluster)samples = np.zeros((n_sample_per_class*n_cluster , dim))for i in range(n_cluster): samples[i*n_sample_per_class: (i+1)*n_sample_per_class] = np.random.randn(n_sample_per_class, dim) + (i+1) * 2 np.savetxt("data.csv", samples, delimiter=",", fmt='%.3f')This code creates 250,000 samples for each cluster, totaling 1 million samples. Each sample has 10 dimensions, and the means of the clusters are 2, 4, 6, and 8 along each dimension, respectively.Rust ImplementationIn this section, I’ll walk you through the Rust implementation of the K-means algorithm step by step. However, since I’m not a seasoned Rust developer, I’ll keep the explanations high-level to avoid potential errors.Setting Up DependenciesThis code snippet imports several standard library and external crates commonly used in Rust for various functionalities. It includes error handling, CSV file reading and writing, array operations, random number generation, and command-line argument parsing.use std::{ error::Error,};use std::ops::IndexMut;use csv::{ReaderBuilder, Writer};use ndarray::{Array1, Array2, ArrayView1, AssignElem};use ndarray_csv::{Array2Reader};use rand::seq::index::sample;use rand::{Rng, thread_rng};use clap::{Parser};Defining Command-Line ArgumentsThis section defines the command-line arguments using the clap crate. It includes arguments for the path to the CSV file, the number of clusters, whether to use Kmeans++ for center initialization, the maximum number of iterations, the tolerance for center change, and the path to save the indices as a CSV file.// define arguments using clap#[derive(Parser)]#[command(author, version, about, long_about = None)]struct Args { #[arg(short, long, help = "Path to the csv file")] data_path: String, #[arg(short, long, help = "Number of clusters")] num_cluster: usize, #[arg(short, long, help = "Use Kmeans++ to initialize centers")] kpp: bool, #[arg(short, long, help = "Maximum number of iterations", default_value = "1000")] max_iter: i32, #[arg(short, long, help = "Maximum center change tolerance", default_value = "1e-4")] tolerance: f32, #[arg(short, long, help = "Path to save indices as csv", default_value = "indices.csv")] output_path: String}Loading Data from CSVThis function load_data takes a file path as input and loads data from a CSV file into a 2D array of 32-bit floating-point numbers (Array2). It uses the csv crate to read the CSV file and deserialize its contents into the array.fn load_data(file_path: &str) -> Result<Array2<f32>, Box<dyn Error>> { // this function loads data from the csv file let reader = ReaderBuilder::new().has_headers(false).from_path(file_path); let array_read: Array2<f32> = reader?.deserialize_array2_dynamic()?; Ok(array_read)}Initializing Cluster CentersThe random_centers function takes a reference to a 2D array data and the number of clusters n_cluster as input. It selects n_cluster random rows from data as initial cluster centers and returns them in a new 2D array.Similarly, the kmeans_pp function implements the KMeans++ algorithm for selecting initial cluster centers.fn random_centers(data : &Array2<f32>, n_cluster: usize) -> Array2<f32> { // centers are randomly selected from the current set of samples let n_rows = data.nrows(); let mut rng = thread_rng(); let mut selected_rows = Array2::<f32>::zeros((n_cluster, data.ncols())); let indices: Array1<usize> = sample(&mut rng, n_rows, n_cluster).into_iter().collect(); for (i, &index) in indices.iter().enumerate() { selected_rows.row_mut(i).assign(&data.row(index)); } selected_rows}fn kmeans_pp(data : &Array2<f32>, n_cluster: usize) -> Array2<f32> { // centers are selected based on KMeans++ algorithm let n_rows = data.nrows(); let mut chosen_points: Vec<usize> = vec!(); let mut centers: Array2<f32> = Array2::zeros((n_cluster, data.ncols())); // first center chosen_points.push(thread_rng().gen_range(0..n_rows)); centers.row_mut(1).assign(&data.row(chosen_points[0])); // other centers for i_center in 1..n_cluster { let mut max_dist: f32 = -1.0; let mut max_index: usize = 0; for (i_sample, sample) in data.rows().into_iter().enumerate() { if chosen_points.contains(&i_sample) { continue } let mut c_dist: f32 = 0.0; for i_prev_centers in 0..i_center { c_dist += euclidean_distance(&sample, &centers.row(i_prev_centers)); } if c_dist > max_dist { max_dist = c_dist; max_index = i_sample } } chosen_points.push(max_index); centers.row_mut(i_center).assign(&data.row(max_index)) } centers}Calculating Euclidean DistanceThe euclidean_distance function calculates the Euclidean distance between two 1D arrays a and b represented as ArrayView1.fn euclidean_distance(a: &ArrayView1<f32>, b: &ArrayView1<f32>) -> f32 { // calculate Euclidean distance between two arrays a.iter().zip(b.iter()).map(|(x, y)| (x - y).powi(2)).sum::<f32>().sqrt()}Assigning Samples to ClustersThe assign_cluster_to_sample function assigns each sample in the input data (Array2) to the closest cluster center based on Euclidean distance. It takes three arguments: the input data, the cluster centers, and a mutable array of indices representing the assigned cluster for each sample.fn assign_cluster_to_sample(data: &Array2<f32>, centers: &Array2<f32>, indices: &mut Array1<usize>) { // finds the closest cluster to each sample using centers for (index_sample, row) in data.rows().into_iter().enumerate(){ let mut min_distance = f32::INFINITY; let mut min_index: usize = 0; for (index, center_row) in centers.rows().into_iter().enumerate(){ let distance = euclidean_distance(&row, &center_row); if distance < min_distance{ min_distance = distance; min_index = index; } } indices.index_mut(index_sample).assign_elem(min_index) }}Updating Cluster CentersThe update_centers function updates the cluster centers based on the samples assigned to each cluster. It takes the input data (data), the current cluster centers (centers), the indices representing the assigned cluster for each sample (indices), and the number of clusters (n_clusters) as input.fn update_centers(data: &Array2<f32>, centers: &mut Array2<f32>, indices: &Array1<usize>, n_clusters: usize) -> f32 { // update centers as the average of the samples within the cluster let mut max_change = f32::INFINITY; for index in 0..n_clusters { let matched_indices: Vec<usize> = indices.iter() .enumerate() .filter(|&(_, &value)| value == index as usize) .map(|(i, _)| i) .collect(); let mut c: Array1<f32> = Array1::zeros(data.ncols()); for m_index in &matched_indices { c += &data.row(*m_index); } c /= matched_indices.len() as f32; let distance = euclidean_distance(&centers.row(index), &c.view()); centers.row_mut(index as usize).assign(&c); if distance < max_change{ max_change = distance; } } max_change}###The write_csv function writes the cluster assignments (indices) to a CSV file specified by output_pathfn write_csv(indices: &Array1<usize>, output_path: &str) -> Result<(), Box<dyn Error>>{ // write the result into a csv file let mut writer = Writer::from_path(output_path)?; for &value in indices.iter() { writer.write_record(&[value.to_string()])?; } writer.flush()?; println!("Results wrote to {output_path}"); Ok(())}K-Means Struct and MethodsThe KMeans struct represents a K-means clustering model with configurable parameters. It includes fields for the number of clusters (num_cluster), whether to use KMeans++ for center initialization (kpp), the convergence tolerance (tolerance), the maximum number of iterations (max_iter), the current iteration number (iter), and the maximum change in cluster centers (max_change).The fit method fits the K-means model to the input data (data). It initializes the cluster centers, assigns samples to clusters, and updates the centers iteratively until convergence or the maximum number of iterations is reached. It returns an array of cluster indices indicating the cluster assignment for each sample.#[derive(Default)]struct KMeans { num_cluster: usize, kpp: bool, tolerance: f32, max_iter: i32, iter: Option<i32>, max_change: Option<f32>}impl KMeans { fn fit(&mut self, data: &Array2<f32>) -> Array1<usize> { // initiate centers random or using kmeans++ let mut centers = if self.kpp { kmeans_pp(&data, self.num_cluster) } else { random_centers(&data, self.num_cluster) }; // main loop, assign indices and update centers let mut indices: Array1<usize> = Array1::zeros(data.nrows()); let mut iter = 0; let mut max_change = f32::INFINITY; while (max_change > self.tolerance) & (iter < self.max_iter) { iter += 1; assign_cluster_to_sample(&data, &centers, &mut indices); max_change = update_centers(&data, &mut centers, &indices, self.num_cluster); } let _ = self.iter.insert(iter); let _ = self.max_change.insert(max_change); indices } fn get_iter(&self) -> i32 { if self.iter.is_some() { self.iter.unwrap() } else { 0 } } fn get_max_change(&self) -> f32 { if self.max_change.is_some() { self.max_change.unwrap() } else { f32::INFINITY } } }Main FunctionThe main function serves as the entry point for the program. It parses command-line arguments using clap, loads data from a CSV file, initializes a KMeans struct with the specified parameters, fits the K-means model to the data, prints the number of iterations and maximum change in cluster centers, and writes the cluster assignments to a CSV file.fn main() { // parse and get arguments let cli = Args::parse(); let file_path = cli.data_path; // load data let data = load_data(&file_path).expect("Error reading csv"); let mut kmeans = KMeans { num_cluster: cli.num_cluster, kpp: cli.kpp, tolerance: cli.tolerance, max_iter: cli.max_iter, ..Default::default() }; let indices = kmeans.fit(&data); println!("Number of Iters: {:?} with max change: {:?}", kmeans.get_iter(), kmeans.get_max_change()); // write to csv file let _ = write_csv(&indices, &cli.output_path);}Usage>> target/release/kmeans -hA simple implementation of KMeans algorithm.Usage: kmeans [OPTIONS] --data-path <DATA_PATH> --num-cluster <NUM_CLUSTER>Options: -d, --data-path <DATA_PATH> Path to the csv file -n, --num-cluster <NUM_CLUSTER> Number of clusters -k, --kpp Use Kmeans++ to initialize centers -m, --max-iter <MAX_ITER> Maximum number of iterations [default: 1000] -t, --tolerance <TOLERANCE> Maximum center change tolerance [default: 1e-4] -o, --output-path <OUTPUT_PATH> Path to save indices as csv [default: indices.csv] -h, --help Print help -V, --version Print versionA Speed Showdown: Rust vs. Python for K-Means ClusteringThis comparison pits the performance of Rust against Python for K-means clustering. The Python code uses numpy and scikit-learn to read data from a CSV file, perform K-means clustering, and save the cluster centers to another CSV file.import timeimport numpy as npfrom sklearn.cluster import KMeansstart = time.time()data = np.genfromtxt('data.csv', delimiter=',')k = KMeans(n_clusters=4, init='k-means++', n_init=1, max_iter=100, tol=1e-5)k.fit(data)np.savetxt("res.csv", k.cluster_centers_, delimiter=",")print(f'n_iter {k.n_iter_}')print(time.time() - start)Despite the fact that I’m not a Rust developer, the Rust code outperforms the Python code by a factor of 7, showcasing Rust’s speed and efficiency.Code RepositoryThe complete Rust implementation of the K-means algorithm is available on GitHub.Reference[1] Rust Book[2] ChatGPT

Gender Bias in Multimodal Embeddings: An Example of OpenAI CLIP

11 minute read

Published: March 03, 2024

Gender Bias in Multimodal Embeddings: An Example of OpenAI CLIP Introduction:Artificial intelligence has made significant strides in understanding and processing multimodal information, such as images and text, simultaneously. OpenAI’s CLIP (Contrastive Language–Image Pretraining) model is a prime example of this advancement, showcasing the ability to learn from diverse image and text data to perform various tasks, including image classification based on textual descriptions. However, as with any AI model, CLIP is not immune to biases present in its training data, raising concerns about its performance and behavior in real-world applications.In this blog post, I investigate the issue of gender bias in multimodal embeddings, focusing on our analysis of OpenAI CLIP using the UTKFace dataset. By examining how CLIP processes and represents gender-related information, I aim to shed light on potential biases and their implications.Our methodology involves comparing embeddings of images paired with gender-related words to predict the gender of the faces in those images. Additionally, I investigate CLIP’s tendency towards associating certain attributes with specific genders, highlighting the challenges and considerations in developing fair and unbiased AI models.Packages and DatasetTo conduct our analysis, I utilized Google Colab for its powerful computational capabilities and convenient access to the necessary libraries. Let’s begin by installing the required packages.!pip -q install ftfy regex tqdm!pip -q install git+https://github.com/openai/CLIP.gitNext, we need to download the UTKFace dataset, which contains a diverse collection of face images labeled with age, gender, and ethnicity information. The dataset is split into three parts, which I downloaded and extracted into a directory named utk_face for further processing.!gdown 1mb5Z24TsnKI3ygNIlX6ZFiwUj0_PmpAW!gdown 19vdaXVRtkP-nyxz1MYwXiFsh_m_OL72b!gdown 1oj9ZWsLV2-k2idoW_nRSrLQLUP3hus3b!mkdir utk_face!tar -xf 'part1.tar.gz' -C './utk_face'!tar -xf 'part2.tar.gz' -C './utk_face'!tar -xf 'part3.tar.gz' -C './utk_face'Creating the DatasetTo create a PyTorch dataset for the UTKFace dataset, I defined a custom UTKFace class that inherits from torch.utils.data.Dataset. This class prepares the samples by loading images and their corresponding labels (age, gender, and race) from the dataset directory. We also included mappings for gender and race labels for easier interpretation.import globimport osfrom torch.utils import dataclass UTKFace(data.Dataset): gender_map = {0: 'male', 1: 'female'} race_map = {0: 'white', 1: 'black', 2: 'asian', 3: 'indian', 4: 'others'} def __init__(self, root, transform=None): self.root = root self.samples = self._prepare_samples(root) self.transform = transform def __getitem__(self, index): path, label = self.samples[index] image = Image.open(path) if self.transform is not None: image = self.transform(image) return image, label def __len__(self): return len(self.samples) def _prepare_samples(self, root): samples = [] paths = glob.glob(os.path.join(root, '*/*')) for path in paths: try: label = self._load_label(path) except Exception as e: print(f'path: {path}, exception: {e}') samples.append((path, label)) return samples def _load_label(self, path): str_list = os.path.basename(path).split('.')[0].strip().split('_') age, gender, race = map(int, str_list[:3]) label = dict(age=age, gender=gender, race=race) return labelTo visualize an example image from the dataset, I randomly selected one image and displayed it using matplotlib.import randomimport matplotlib.pyplot as pltfrom PIL import Imageutkface = UTKFace(root='utk_face')print(f'num images: {len(utkface)}')sample_image, sample_label= random.choice(utkface)plt.imshow(sample_image)plt.axis("off")plt.title(str(sample_label))plt.show()Model and DataLoader SetupWe selected the CLIP model architecture “ViT-B/32” for our analysis, which is a Vision Transformer model with a patch size of 32x32 pixels. We also set up the data loader to efficiently load and preprocess images from the UTKFace dataset for inference.import torchimport clipfrom PIL import Imagefrom tqdm import tqdmprint('Available Models: ', clip.available_models())device = "cuda" if torch.cuda.is_available() else "cpu"print(f'device is {device}')model, preprocess = clip.load("ViT-B/32", device=device)model.eval()utkface_dataset = UTKFace(root='utk_face', transform=preprocess)utkface_dataloader = data.DataLoader(utkface_dataset, batch_size=512, shuffle=False, num_workers=2)Gender PredictionHere’s the code for predicting gender based on the provided keywords:import numpy as np# change this to use one or more keywordstext_tokens = []text_tokens.append(clip.tokenize(["male", "female"]).to(device))# text_tokens.append(clip.tokenize(["man", "woman"]).to(device))# text_tokens.append(clip.tokenize(["boy", "girl"]).to(device))n = len(text_tokens)age = []gender = []race = []gender_p = []with torch.no_grad(): text_feats = [] for images, label in tqdm(utkface_dataloader): age.extend(label["age"].tolist()) gender.extend(label["gender"].tolist()) race.extend(label["race"].tolist()) probability_g = np.zeros((images.shape[0], 2)) for tt in text_tokens: logits_per_image, logits_per_text = model(images.to(device), tt) probs = logits_per_image.softmax(dim=-1).cpu().numpy() probability_g += probs / n gender_p.extend(probability_g[:,1])This code calculates the gender probabilities for each image in the dataset based on the provided keywords. The gender_p list will contain the predicted probabilities of being female for each image. You can adjust the text_tokens list to include different sets of keywords for gender prediction.Generating Ground Truth DataTo ensure the integrity of our analysis, let’s filter out samples with invalid gender or race labels from the UTKFace dataset. I retained only samples where the gender label was less than 3 (indicating male or female) and the race label was less than 5 (indicating one of the specified races). This filtering process resulted in a total of 24,105 samples for which we have both ground truth gender labels and predicted gender probabilities.age = np.array(age)gender = np.array(gender)race = np.array(race)gender_p = np.array(gender_p)ind_keep = (gender < 3) & (race < 5)age = age[ind_keep]gender = gender[ind_keep]race = race[ind_keep]gender_p = gender_p[ind_keep]print(f'Shape of ground truth is {gender_p.shape}')Optimizing Threshold and Calculating MetricsTo evaluate the performance of our gender prediction model based on the predicted gender probabilities from OpenAI CLIP, I first classified the probabilities using a threshold of 0.5 to determine the predicted gender class. We then optimized the threshold to maximize the F-score, a metric that balances precision and recall, for gender prediction.from sklearn.metrics import precision_recall_fscore_supportfrom scipy import statsgender_p_class = (gender_p > 0.5).astype(int)mode = stats.mode(gender_p_class)mode_class = UTKFace.gender_map[mode.mode]print(f'Gender has bias towards {mode_class}.')print(f'Predicted #{mode_class}: {np.mean(gender_p_class==mode.mode)*100:.2f}%/ True value of #{mode_class}: {np.mean(gender==mode.mode)*100:.2f}%')# optimize thresholdfmax = 0threshold = Nonefor th in np.arange(0.05, 1, 0.05): p, r, f, s= precision_recall_fscore_support(gender, gender_p>th) if np.mean(f) > fmax: fmax = np.mean(f) threshold = thprint(f'\nF-score max: {fmax:.4f} with threshold: {threshold}')print(f'Accuracy (Threshold: {threshold}): {np.mean((gender_p > threshold) == gender)}')print(f'Accuracy (Threshold: {0.5}): {np.mean((gender_p > 0.5) == gender)}')This code snippet provides insights into the bias present in the predicted gender classes and demonstrates the optimization of the threshold for gender prediction. The resulting F-score and accuracy metrics help evaluate the performance of the gender prediction model and provide a basis for further analysis.Error Analysis by RaceTo understand how gender prediction errors vary across different races in the dataset, I calculated the error rate for each race. The error rate is defined as the proportion of samples for which the predicted gender class differs from the ground truth gender class.gender_p_class = (gender_p > threshold).astype(int)race_error = race[gender != gender_p_class]unique, counts = np.unique(race_error, return_counts=True)for u, c in zip(unique, counts): print(f'{UTKFace.race_map[u]} : {c/len(race_error):.4f}')This analysis provides insights into whether the gender prediction model exhibits biases in predicting gender across different racial groups. The results can help identify potential areas for improvement and ensure the fairness of the model across diverse populations. To further analyze the gender prediction performance across different racial groups, I computed confusion matrices for each race. A confusion matrix provides a detailed breakdown of correct and incorrect predictions made by the gender prediction model.from sklearn.metrics import confusion_matrixgender_p_class = (gender_p > threshold).astype(int)for ur in np.unique(race): print(UTKFace.race_map[ur]) print(confusion_matrix(gender[race==ur], gender_p_class[race==ur])) print('\n')Results: Analysis of Gender Bias in CLIPMale vs. FemaleWhen using the keywords “male” and “female” to predict gender, the model exhibited a bias towards predicting male gender. The predicted percentage of males in the dataset was 66.30%, compared to the true value of 52.20%. This bias highlights the need to carefully select keywords and thresholds for gender prediction.By optimizing the threshold to 0.2, we achieved a maximum F-score of 0.9603. This threshold indicates that for an image to be classified as female, the probability of being female should be at least 0.2 which shows the bias towards male. At this threshold, the accuracy was 0.9604, significantly higher than the accuracy at the default threshold of 0.5 (0.8503).The error rates for gender prediction varied across different racial groups. For the white race, the error rate was 0.3309, indicating a relatively high rate of misclassification. In contrast, the error rates for black, asian, indian, and other races were lower, ranging from 0.1173 to 0.2785. These results suggest that the model’s performance in predicting gender may be influenced by racial factors.Man vs. WomanUsing the keywords “man” and “woman” to predict gender, the model exhibited a bias towards predicting male gender. The predicted percentage of males was 53.65%, compared to the true value of 52.20%. This bias underscores the importance of carefully selecting keywords and thresholds for gender prediction to mitigate biases.By optimizing the threshold to 0.3, we achieved a maximum F-score of 0.9487. This threshold indicates that for an image to be classified as female, the probability of being female should be at least 0.3. At this threshold, the accuracy was 0.9487, which is slightly higher than the accuracy at the default threshold of 0.5 (0.9421).The error rates for gender prediction varied across different racial groups. The error rate was highest for the asian race at 0.3641, followed by the white race at 0.2921. The error rates for black, indian, and other races were lower, ranging from 0.0963 to 0.1472. These results suggest that the model’s performance in predicting gender may be influenced by racial factors.Boy vs. GirlUsing the keywords “boy” and “girl” to predict gender, the model exhibited a bias towards predicting male gender. The predicted percentage of males was 61.34%, compared to the true value of 52.20%.By optimizing the threshold to 0.2, we achieved a maximum F-score of 0.9412. This threshold indicates that for an image to be classified as female, the probability of being female should be at least 0.2, indicating a high bias towards male. At this threshold, the accuracy was 0.9413, which is higher than the accuracy at the default threshold of 0.5 (0.8865).The error rates for gender prediction varied across different racial groups. The error rate was highest for the white race at 0.4064, followed by the asian race at 0.2615. The error rates for black, indian, and other races were lower, ranging from 0.0777 to 0.1569.ColculsionIn this blog post, we analyzed gender bias in OpenAI CLIP using different keyword pairs (“male” and “female”, “man” and “woman”, “boy” and “girl”) to predict gender in images. We investigated the model’s bias, optimized thresholds for gender prediction, and evaluated its performance across different racial groups. Gender Bias: Across all keyword pairs, the model exhibited a bias towards predicting male gender. The predicted percentages of males were consistently higher than the true values, indicating a systematic bias in the model’s predictions. Optimized Thresholds: By optimizing the threshold for gender prediction, we were able to achieve higher accuracies and F-scores compared to the default threshold of 0.5. This suggests that adjusting the threshold can improve the model’s performance and reduce bias. Error Analysis by Race: The error rates for gender prediction varied across different racial groups. The model performed relatively poorly for the white and asian races, indicating potential challenges in accurately predicting gender for these groups. What We Learned In multimodal models, the choice of keywords and thresholds is crucial in mitigating bias and improving the performance of gender prediction models. Racial diversity in the dataset can impact the model’s performance, highlighting the importance of diverse and representative datasets.ColabYou can find the code at Colab.Reference[1] CLIP[2] UTKFace[2] ChatGPT

Weak Supervision with Snorkel: Image Classification Example

13 minute read

Published: August 29, 2023

Weak Supervision with Snorkel: Image Classification Example Introduction:In the world of machine learning, data is often hailed as the crown jewel that powers models and drives innovation. Yet, obtaining high-quality, labeled data remains a significant challenge, often demanding painstaking manual efforts from human annotators. This is where the concept of weak supervision emerges as a beacon of hope for machine learning engineers and practitioners.Weak supervision is the art of leveraging various sources of noisy or imprecise supervision to label a large amount of data efficiently. It takes the burden off exhaustive manual labeling and opens the door to scaling up projects that might have been otherwise resource-intensive. In this post, we embark on a journey to explore the Snorkel, a powerful tool that empowers us to automate the labeling process, saving time and effort without compromising on results.In this tutorial, tailored for machine learning engineers and enthusiasts alike, we’ll unveil the advantages of weak supervision using a practical example: Image Classification. By the end of this guide, you’ll have a basic understanding of how to harness the potential of Snorkel to streamline your image classification pipelines and achieve impressive results with reduced labeling efforts.Whether you’re a seasoned practitioner seeking to optimize your workflow or a newcomer eager to unlock the potential of weak supervision, this tutorial will equip you with the knowledge and skills needed to elevate your machine learning projects. So, let’s dive into the world of weak supervision and see how Snorkel can revolutionize the way we approach labeling and ultimately, supercharging our machine learning models.Are you ready to embark on this exciting journey? Let’s begin!Data Download: Exploring the Open Images Dataset V7Before we dive into the exciting world of weak supervision and Snorkel for image classification, we need to set the stage by obtaining the necessary data. In this tutorial, we’ll be using the Open Images Dataset V7, a rich collection of images spanning a wide array of categories. This dataset is a treasure trove for machine learning tasks, providing a diverse range of visuals that will help us showcase the power of weak supervision.To get started, we’ll perform a series of commands to download the essential files from the Open Images Dataset V7. These files contain crucial information about class labels, class descriptions, and annotations. Below is the code snippet you’ll need to execute to gather these files:In this set of commands, we create a directory named oiv7 to neatly organize the downloaded files. The downloaded files include: oidv7-classes-trainable.txt: A list of trainable (verified) class labels. oidv7-class-descriptions.csv: A CSV file containing class descriptions. oidv7-train-annotations-human-imagelabels.csv: Annotations for training images. oidv7-val-annotations-human-imagelabels.csv: Annotations for validation images. oidv7-test-annotations-human-imagelabels.csv: Annotations for test images.Now that we have the essential data files in place, it’s time to turn our attention to the actual image files. In this section, we’ll walk through the process of downloading labeled images that are trainable according to the Open Images Dataset V7.The provided Python code streamlines this image download process, ensuring that we only retrieve images that are relevant and fit for training.Let’s break down the key components of the code: download_one_image: This function downloads a single image from the specified split (train, val, or test) and saves it to the specified path. The function uses the BUCKET resource from the boto3 library to interact with the S3 bucket. get_class_label: This function retrieves class labels for the requested class names. It ensures that the requested classes are trainable, as per the dataset specifications. get_image_ids: This function retrieves image IDs for requested splits and labels. It identifies images that match the requested class labels and have a confidence level of 1.0 (verified by human). download_images: This function orchestrates the image download process based on requested labels, splits, and paths. It uses concurrent futures to speed up the download process by using multiple threads. By combining these components, the code provides a way to download labeled images from the Open Images Dataset V7. Next, we’ll use the code to download labeled images,specifically focusing on the classes “sea” and “Jungle.” These images will serve as our starting point for weak supervision, demonstrating how we can leverage Snorkel to automatically label and train an image classifier.Organizing and Splitting the DatasetIn this section, we’ll walk through a code snippet that performs data splitting and organization. This step is pivotal in setting the stage for robust model training and evaluation.Let’s delve into the key components of the code: stat: This function computes statistics on the data by counting the number of images for each class. It returns a dictionary mapping class names to lists of image paths. split_manual: This function splits the data into different subsets (train, val, test) based on specified ratios. It ensures that each subset maintains a proportional representation of different classes. We split the data into train, val and test proportional to 0.7, 0.2, and 0.1. Labeling FunctionsThe heart of the Snorkel lies in creating labeling functions that generate noisy labels for our data.In this section, we’ll go though a code script that defines a set of labeling functions, each contributing to the creation of our weakly labeled dataset.Here’s an overview of the key components of the provided code: check_color: This function classifies images based on their dominant color on average. check_pixel_color: This function classifies images based on the mode of max color per pixel. check_with_efficientNet: This function leverages EfficientNet predictions and FastText embeddings to classify images as “SEA” or “JUNGLE.” Here, first we classify the image using EfficientNet. Then, the closeness of the output label is examinedagainst several words related to sea or jungle with FastText. Thus, if the meaning of the output labelis closer to sea, we label it as SEA. Adding labeling functions: The script uses the add_func decorator to add each labeling function to the LABELING_FUNCS list, which will be used later. By combining these labeling functions, we will generate a set of noisy labels for our images. These labels are the cornerstone of our weak supervision approach, allowing us to utilize Snorkel’s capabilities.Weak Supervision Labeling with SnorkelThe power of weak supervision comes to life when we leverage labeling functions to create noisy labels for our dataset. Now, we’ll explore a script that performs weak supervision labeling with Snorkel.Here’s a breakdown of the key elements in the provided code: DATA PREPARATION: The script starts by specifying the root directory containing the input images and the splits (e.g., ‘train’, ‘val’) to process. Additionally, it defines the root directory to save the labeled data. LFApplier: The labeling functions (LABELING_FUNCS) defined in the previous code script are applied to all images in the specified split. The result is a label matrix (L_train) where each row corresponds to an image and each column corresponds to a labeling function. LFAnalysis: This step provides an analysis of the labeling functions’ performance on the data. It generates a summary that indicates how well the labeling functions agree or disagree on assigning labels to images. label_model: A LabelModel is trained using the label matrix (L_train). This model learns to estimate the true underlying labels by accounting for the noise introduced by the labeling functions. Label Prediction: The label model predicts probabilities of labels for each image based on the noisy labels from the labeling functions. Saving Labeled Data: The labeled data, including the predicted labels and image paths, is saved to pickle files. This data will serve as the input for our model training process. By executing this script, we perform the crucial step of labeling our data using weak supervision techniques. Snorkel helps us manage the uncertainty introduced by the labeling functions, creating a labeled dataset that reflects the inherent noise in the weakly supervised data.DataloadersNext, we implement our dataloaders for the supervised and the weakly supervised procedures.Here’s a breakdown of the key elements in the provided code: get_transforms(): This function provides data transformation pipelines tailored for different dataset splits: ‘train’, ‘val’, and ‘test’. These transformations include resizing, cropping, flipping, rotation, normalization, and tensor conversion. SnorkelDataset: This custom dataset class is designed for weakly supervised learning using Snorkel labels. It takes a path to a pickled data file and a label type (‘hard’ or ‘soft’) as inputs. In weakly supervised learning, “hard labels” refer to discrete, definite labels assigned to data points, indicating clear categorization (e.g., ‘SEA’ or ‘JUNGLE’). On the other hand, “soft labels” represent probabilistic or continuous assignments, reflecting the uncertainty or ambiguity in classification. The class loads images and corresponding labels from the data file and applies the specified transformations. get_data_loader(): This function creates and returns data loaders for different dataset splits. It utilizes the ImageFolder dataset from PyTorch, which organizes data into class folders. The dataloaders are configured with appropriate transformations and batch sizes for training, validation, and testing. get_data_loader_snorkel(): This function generates data loaders for the specified dataset splits using Snorkel-generated labels. It utilizes the SnorkelDataset class to load images and labels from pickled data files, enabling weakly supervised learning. The dataloaders are configured similarly to those in get_data_loader(), tailored for Snorkel-labeled data. By leveraging these utility functions and classes, we ensure that our data is well-prepared and ready to be fed into our CNN model.Training functionHere is the training function. I skip the description of this part, since it follows a common pattern of model training with pytorch. NOTE: The condition if len(labels.shape) > 1: deals with soft labels.InferenceWith the following code, we can evaluate all weights saved in the weight_folder on the test data.ResultsNow let’s see the results. With the following script, we can train the model with original supervised procedure or two semi-supervised ones (soft and hard labels).The result isOriginal Labels: Precision: For the original labels, the model achieves high precision scores for both classes (‘SEA’ and ‘JUNGLE’), indicating that when it predicts a class, it’s usually correct. Specifically, the precision values of approximately 0.94 and 0.92 for ‘SEA’ and ‘JUNGLE’ respectively demonstrate the model’s accuracy in its predictions. Accuracy (ACC): The overall accuracy of 0.926 suggests that the model is successful in correctly classifying approximately 92.6% of the images in the dataset. Snorkel Soft Labels: Precision: With soft labels, where the labeling functions provide probabilistic or continuous assignments, the precision scores remain relatively high but show a slight decrease compared to the original labels. The values of around 0.93 for ‘SEA’ and 0.909 for ‘JUNGLE’ indicate a minor decrease in precision. Accuracy (ACC): The accuracy of 0.9132, while slightly lower than the original labels, still demonstrates a strong performance, capturing approximately 91.3% of the dataset correctly.Snorkel Hard Labels: Precision: When using hard labels (discrete, definite labels) provided by Snorkel, there is a more noticeable decrease in precision for the ‘SEA’ class, dropping to approximately 0.882. However, the precision for ‘JUNGLE’ remains high at around 0.923. Accuracy (ACC): The overall accuracy of 0.9144, although slightly lower than the original labels, showcases the model’s ability to maintain a strong classification performance with Snorkel hard labels.ConclusionIn the pursuit of harnessing the power of weak supervision, our journey has traversed a landscape where precision meets ambiguity and accuracy coexists with uncertainty. The application of labeling functions in the Snorkel framework has enabled us to embrace our prior knowledge in our data, offering a nuanced perspective on image classification. We also observed the advantage of using soft labels: their ability to mitigate the impact of class imbalances. In our dataset, where ‘SEA’ and ‘JUNGLE’ classes exhibited varying instances, soft labels allowed for a more nuanced representation of uncertainty. This nuanced understanding ensured that the precision for both classes stayed relatively close, compared to the more discrete hard labels.In an imbalanced dataset, the imprecision introduced by labeling functions might disproportionately affect the minority class. Soft labels, by representing class assignments probabilistically, provided a flexibility that allowed the model to balance the precision between the classes more effectively. This balancing act is crucial, especially in applications where misclassifying the minority class carries significant consequences.It’s important to note that while our labeling functions in this example were relatively straightforward to implement, many real-world scenarios pose complex challenges. Designing accurate labeling functions can be intricate, requiring domain expertise and careful consideration. Despite these challenges, our study demonstrates that noisy labels, when harnessed intelligently through weak supervision techniques, can still offer valuable insights and contribute to robust model training.Our exploration underscores the resilience of machine learning models in the face of noisy or uncertain labels. Even when labeling functions are not perfect, the intelligent integration of these noisy annotations can lead to significant advancements in model performance. Embracing the inherent noise in weakly supervised data and leveraging techniques like Snorkel not only expands the scope of feasible applications but also highlights the adaptability and learning potential of modern machine learning systems.GitHubYou can find the code at the project GitHub repository.Reference[1] Snorkel[2] ChatGPT[3] Open Images Downloader

Optimizing Python Code with LRU Cache: A Fibonacci Sequence Example

2 minute read

Published: May 03, 2023

Optimizing Python Code with LRU Cache: A Fibonacci Sequence Example LRU Cache is a powerful feature in Python that can help optimize the performance of code that involves frequent function calls. LRU stands for “Least Recently Used”, and it is a type of cache that stores the results of recently called functions. It is particularly useful in situations where the same function is called repeatedly with the same arguments, as it can significantly reduce the number of function calls required. LRU Cache is a powerful feature in Python that can help optimize the performance of code that involves frequent function calls. LRU stands for “Least Recently Used”, and it is a type of cache that stores the results of recently called functions. It is particularly useful in situations where the same function is called repeatedly with the same arguments, as it can significantly reduce the number of function calls required.One classic example of a problem that benefits from LRU caching is the Fibonacci sequence. The Fibonacci sequence is a series of numbers in which each number is the sum of the two preceding numbers, starting from 0 and 1. The sequence can be defined recursively as follows:While this implementation is simple and easy to understand, it can quickly become inefficient as the size of n grows. This is because the function calls itself twice for every recursive call, resulting in an exponential increase in the number of function calls. For example, calling fib(6) would result in the following function calls:As you can see, the function is called multiple times with the same arguments, resulting in redundant calculations. This is where LRU caching comes in handy.Python’s functools module provides a built-in decorator lru_cache that can be used to cache the results of a function. Here’s how we can use it to optimize the Fibonacci sequence:By adding the @lru_cache decorator to the function, we tell Python to cache the results of previous function calls. The maxsize argument sets the maximum number of function calls that can be stored in the cache. Setting it to None means that the cache can store an unlimited number of function calls.Let’s now test our optimized implementation of the Fibonacci sequence by calling fib(6) again:As you can see, the number of function calls has been significantly reduced, and the function only calculates each value once. This makes our implementation much more efficient and faster, especially for larger values of n.In this post we learned that LRU caching is a powerful optimization technique that can significantly improve the performance of Python code that involves frequent function calls. By using the @lru_cache decorator, we can easily implement caching in our code and reduce the number of redundant function calls. The Fibonacci sequence is just one example of a problem that can benefit from LRU caching, but there are many other use cases where caching can make a big difference in performance.Reference[1] Fluent Python

A += Assignment Puzzler in Python

1 minute read

Published: April 05, 2023

A += Assignment Puzzler in Python In the world of programming, it’s not uncommon to come across some puzzling situations that require careful analysis and attention to detail. This is exactly the case with this post, which involves evaluating the following code. Let’s dive into this puzzle and see what we can learn from it. In the world of programming, it’s not uncommon to come across some puzzling situations that require careful analysis and attention to detail. This is exactly the case with this post, which involves evaluating the following code. Let’s dive into this puzzle and see what we can learn from it.What happens next? Choose the best answer: t becomes (1, 2, [30, 40, 50, 60]). TypeError is raised with the message ‘tuple’ object does not support item assignment. Neither. Both 1 and 2.Tuples are immutable and cannot be changed. Therefore, if we try to change an element of a tuple, a TypeError should be raised. However, tuples store the reference to their members and a list can be changed without changing its reference. Thus, you can modify a list in the tuple! The answer to the previous question is actually 4, “Both 1 and 2”. Here is the output:As you can see, we have a TypeError states that “‘tuple’ object does not support item assignment”, and surprisingly, t is also changed.What happened?Here is what happens in Python: Python puts the value of s[2] on TOS (Top Of Stack). Perform TOS += b. This succeeds since TOS refers to a mutable object in our example (a list). Assign s[2] = TOS. This fails since s is immutable in our example (a tuple).We take two lessons from this: Avoid putting mutable items in tuples. Augmented assignment is not an atomic operation—we just saw it throwing an exception after doing part of its job.Reference[1] Fluent Python

Python Trick: Merging Two Lists with the zip Function

1 minute read

Published: March 05, 2023

Python Trick: Merging Two Lists with the zip Function As a data scientist or machine learning engineer, you’ll often need to combine or merge two lists of data. While there are several ways to achieve this in Python, using the zip function is an efficient and concise method that you should have in your code. As a data scientist or machine learning engineer, you’ll often need to combine or merge two lists of data. While there are several ways to achieve this in Python, using the zip function is an efficient and concise method that you should have in your code.The zip function takes two or more iterable objects as arguments and returns a new iterator that aggregates elements from each iterable. This makes it an ideal candidate for merging two lists.Here’s an example:In this example, the zip function merges the two lists into a list of tuples, where each tuple contains an element from both lists.You can also use the zip function to merge more than two lists:Note that the zip function returns an iterator, so you need to convert it to a list if you want to access the merged data multiple times.You can also use the zip function with other iterable objects, such as sets and strings:Using the zip function to merge two or more lists is a Python trick that can save you time and make your code more efficient. Keep it in mind the next time you need to combine data from different sources.

Lazy Predict: Explore lots of machine learning models at once

less than 1 minute read

Published: February 21, 2023

Lazy Predict: Explore lots of machine learning models at once The Lazy Predict module in Python is a library that allows users to quickly and easily create predictive models from data. It provides a simple API for creating and training models, as well as a set of tools for evaluating and optimizing them. An example of using the Lazy Predict module would be to create a model to predict the price of a house based on its size, location, and other features. The user would first create a model using the Lazy Predict API, then train it on a dataset of house prices. Finally, the user would use the tools provided by the module to evaluate and optimize the model, and then use it to make predictions on new data.ExampleHere is a simple example of how Lazy Predict can be used: The Lazy Predict module in Python is a library that allows users to quickly and easily create predictive models from data. It provides a simple API for creating and training models, as well as a set of tools for evaluating and optimizing them. An example of using the Lazy Predict module would be to create a model to predict the price of a house based on its size, location, and other features. The user would first create a model using the Lazy Predict API, then train it on a dataset of house prices. Finally, the user would use the tools provided by the module to evaluate and optimize the model, and then use it to make predictions on new data.ExampleHere is a simple example of how Lazy Predict can be used:Considering the breast cancer dataset, let’s test lots of models on this dataset with Lazy Predict.And here is the output:Reference[1] Lazy Predict

Accelerating Python code with Numba

1 minute read

Published: February 12, 2023

Accelerating Python code with Numba Numba is an open-source numerical Python compiler that translates a subset of Python and NumPy code into machine code, executed on the CPU. It can be used to significantly speed up the execution of numerical computations, especially those that are performed in a loop, by compiling the code into machine code rather than interpreting it dynamically. Numba is an open-source numerical Python compiler that translates a subset of Python and NumPy code into machine code, executed on the CPU. It can be used to significantly speed up the execution of numerical computations, especially those that are performed in a loop, by compiling the code into machine code rather than interpreting it dynamically.Numba provides a just-in-time (JIT) compiler, which means that it compiles the code on-the-fly, just before it is executed. This allows for the optimizations to be applied exactly where they are needed, rather than in advance. Numba supports a wide range of Python data types and functions, and it provides a simple, easy-to-use API for adding JIT compilation to existing code.Numba allows you to write high-performance Python code by leveraging the power of JIT compilation. It can be used for a wide range of applications, from scientific computing and data analysis to machine learning.ExampleHere is a simple example of how Numba can be used to speed up a calculation:Consider the following function:This function can be slow when working with large inputs. To speed up the calculation, we can use Numba to compile the function:By adding the @jit decorator, Numba compiles the function into machine code, which can be executed much faster than the interpreted code. The compiled function can then be used just like any other Python function:The output of the above script is:In this example, we see a significant speedup in the execution time of the function. The same technique can be used to speed up other types of calculations, such as matrix operations, Monte Carlo simulations, and more.Reference[1] A ~5 minute guide to Numba

What is Artificial Intelligence?

3 minute read

Published: February 07, 2023

What is Artificial Intelligence? Artificial Intelligence (AI) is a fascinating and rapidly growing field that has the potential to revolutionize the world as we know it. Simply put, AI refers to the development of computer systems that can perform tasks that normally require human intelligence. This includes everything from understanding natural language, recognizing objects in images, and making decisions to solving complex problems. Artificial Intelligence (AI) is a fascinating and rapidly growing field that has the potential to revolutionize the world as we know it. Simply put, AI refers to the development of computer systems that can perform tasks that normally require human intelligence. This includes everything from understanding natural language, recognizing objects in images, and making decisions to solving complex problems.The Core Technologies in AIAt its core, the goal of AI is to create machines that can think and reason like humans. To achieve this, AI systems are divided into two categories: narrow or weak AI and general or strong AI. Narrow AI is designed to perform specific tasks, such as playing a game or identifying objects in images, while General AI has the potential to perform any intellectual task that a human can.One of the key technologies used in AI is machine learning, which allows computers to learn from data and improve their performance over time. Machine learning algorithms can be supervised, unsupervised, or reinforcement-based, each of which offers unique benefits. For example, supervised learning involves training a model on a labeled dataset with the correct output provided for each input. Unsupervised learning, on the other hand, trains a model on an unlabeled dataset and allows the model to find patterns and relationships on its own. Reinforcement learning involves trial and error training, where the model receives rewards or punishments based on its actions.Another crucial technology in AI is deep learning, which is a subfield of machine learning that uses neural networks to model complex patterns in data. Neural networks are composed of multiple layers of interconnected nodes, and they can be trained on large amounts of data to make predictions or perform tasks such as image classification and natural language processing.Impact of AI on Different IndustriesThe impact of AI is already being felt in various industries, including healthcare, finance, transportation, and education. In healthcare, AI is being used to diagnose diseases, predict patient outcomes, and develop personalized treatment plans. In finance, AI is being used to detect fraud, analyze market trends, and make investment decisions. In transportation, AI is being used to develop self-driving cars and optimize routing and scheduling. And in education, AI is being used to personalize learning experiences and provide real-time feedback to students.Ethical and Social Concerns of AIHowever, the development of AI also raises important ethical and social concerns. For example, the use of AI in the workplace may result in job loss, as machines can perform tasks more efficiently than humans. Additionally, the use of AI in decision-making may result in discrimination and bias, as AI systems may perpetuate and amplify existing biases in the data they are trained on.To ensure that AI is developed and used responsibly, it is essential to consider the ethical and social implications of AI and regulate its development and use. This can be achieved through the development of ethical frameworks and guidelines, as well as through collaboration between researchers, policymakers, and industry.ConclusionAI is a field with tremendous potential to change the world for the better. However, it is critical to approach AI with caution and consider its impact on society as a whole. By collaborating and regulating the development and use of AI, we can create a future where AI is used for the benefit of all.[Generated By ChatGPT]

projects

Real-Time Face Search in Video [Read More]

Published: January 01, 2018

Human Pose Estimation and Motion Capture [Read More]

Published: June 01, 2018

Movie Genre Detection [Read More]

Published: January 01, 2019

YouTube Downloader [Read More]

Published: January 01, 2019

Twitter Job Title Prediction [Read More]

Published: January 01, 2019

Video Content Creation and Curation [Read More]

Published: December 01, 2019

Robust Offline Spike Sorter (ROSS) [Read More]

Published: January 01, 2020

Mousi — Mouse Tracking and Behavioural Analysis [Read More]

Published: July 08, 2021

Cage Camera for Rodent Behavioural Monitoring [Read More]

Published: August 12, 2021

Brain-Inspired Spatial-Frequency-Aware Networks [Read More]

Published: February 13, 2023

Weak Supervision with Snorkel [Read More]

Published: October 17, 2023

K-Means Clustering in Rust [Read More]

Published: March 06, 2024

Retrieval-Augmented Generation for LLM Systems [Read More]

Published: May 01, 2024

CNN Training & Inference in Rust [Read More]

Published: June 29, 2024

Model Optimization: Quantization and Pruning [Read More]

Published: October 01, 2024

Llama Chatbot in Rust (llama.cpp) [Read More]

Published: October 07, 2024

Diffusion Models from Scratch [Read More]

Published: November 27, 2024

Deep Learning for MRI Reconstruction [Read More]

Published: January 01, 2025

Spatial Frequency Representation in the Inferior Temporal Cortex [Read More]

Published: January 08, 2025

Fine-Tuning LLaMA 3.2-Vision for Product Descriptions [Read More]

Published: February 23, 2025

Multimodal Book Genre Prediction [Read More]

Published: July 01, 2025

Persian Speech Recognition with Whisper [Read More]

Published: July 09, 2025

Regex Engine & grep-Style Search (Rust) [Read More]

Published: December 10, 2025

publications

Improved ensemble growing method for steganalysis of digital media [Read More]

Ramin Toosi, Sadaf Salehkalaibar, Mohammad Ali Akhaee, Multimedia Tools and Applications, 2018 [Link]

Robust image watermarking using sample area quantization [Read More]

Ramin Toosi, Mohammadreza Sadeghi, Mohammad Ali Akhaee, Multimedia Tools and Applications, 2019 [Link]

Blind gain invariant image watermarking using random projection approach [Read More]

Mohammadreza Sadeghi, Ramin Toosi, Mohammad Ali Akhaee, Signal Processing, 2019 [Link]

Time–frequency analysis of keystroke dynamics for user authentication [Read More]

Ramin Toosi, Mohammad Ali Akhaee, Future Generation Computer Systems, 2021 [Link]

Soccer Event Detection Using Deep Learning [Read More]

Ali Karimi, Ramin Toosi, Mohammad Ali Akhaee, arXiv, 2021 [Link]

Fast and Temporal Consistent Video Style Transfer [Read More]

Ali Abbasi, Ramin Toosi, Mohammad Ali Akhaee, International Conference on Pattern Recognition and Image Analysis (IPRIA), 2021 [Link]

An automatic spike sorting algorithm based on adaptive spike detection and a mixture of skew-t distributions [Read More]

Ramin Toosi, Mohammad Ali Akhaee, Mohammad-Reza A Dehaqani, Scientific Reports, 2021 [Link] [PDF] [Code]

Optimum Group Pixel Matching Strategies for Image Steganography [Read More]

Alireza Shahanaghi, Mohammad Ali Akhaee, Saeed Sarreshtedari, Ramin Toosi, International ISC Conference on Information Security and Cryptology (ISCISC), 2021 [Link] [PDF]

The impact of spatial frequency on hierarchical category representation in macaque temporal cortex [Read More]

Esmaeil Farhang, Ramin Toosi, Behnam Karami, Roxana Koushki, Ehsan Rezayat, Farideh Shakerian, Jalaledin Noroozi, Mohammad-Reza A Dehaqani, Nature Communications Biology, 2021 [Link] [PDF]

Listening to Sounds of Silence for Audio replay attack detection [Read More]

Mohammad Hajipour, Mohammad Ali Akhaee, Ramin Toosi, International Conference on Signal Processing and Intelligent Systems (ICSPIS), 2021 [Link]

Job Title Prediction from Tweets Using Word Embedding and Deep Neural Networks [Read More]

Shayan Vassef, Ramin Toosi, Mohammad Ali Akhaee, International Conference on Electrical Engineering (ICEE), 2022 [Link] [Code]

Multinomial Emoji Prediction Using Deep Bidirectional Transformers and Topic Modeling [Read More]

Zahra Ebrahimian, Ramin Toosi, Mohammad Ali Akhaee, International Conference on Electrical Engineering (ICEE), 2022 [Link]

An open-set framework for underwater image classification using autoencoders [Read More]

Azim Akhtarshenas, Ramin Toosi, SN Applied Sciences, 2022 [Link] [PDF]

Multimodal movie genre classification using recurrent neural network [Read More]

Tina Behrouzi, Ramin Toosi, Mohammad Ali Akhaee, Multimedia Tools and Applications, 2022 [Link] [PDF] [Code]

Fast and accurate spectral clustering via augmented Lagrangian [Read More]

Ramin Toosi, Mohammadreza Sadeghi, Hossein B Yazdi, Mohammad Ali Akhaee, Journal of Computational Science, 2022 [Link]

Brain-inspired feedback for spatial frequency aware artificial networks [Read More]

Ramin Toosi, Mohammad Ali Akhaee, Mohammad-Reza A Dehaqani, 2022 56th Asilomar Conference on Signals, Systems, and Computers, 2022 [Link]

Automated Person Identification from Hand Images using Hierarchical Vision Transformer Network [Read More]

Zahra Ebrahimian, Seyed Ali Mirsharji, Ramin Toosi, Mohammad Ali Akhaee, International Conference on Computer and Knowledge Engineering (ICCKE), 2022 [Link]

Soccer Video Event Detection Using Metric Learning [Read More]

Ali Karimi, Ramin Toosi, Mohammad Ali Akhaee, International Conference on Computer and Knowledge Engineering (ICCKE), 2022 [Link]

Hate Sentiment Recognition System For Persian Language [Read More]

Pegah Shams Jey, Arash Hemmati, Ramin Toosi, Mohammad Ali Akhaee, International Conference on Computer and Knowledge Engineering (ICCKE), 2022 [Link]

Bilingual COVID-19 Fake News Detection Based on LDA Topic Modeling and BERT Transformer [Read More]

Pouria Omrani, Zahra Ebrahimian, Ramin Toosi, Mohammad Ali Akhaee, 2023 6th International Conference on Pattern Recognition and Image Analysis (IPRIA), 2023 [Link]

Persian Ezafeh Recognition using Transformer-Based Models [Read More]

Ali Ansari, Zahra Ebrahimian, Ramin Toosi, Mohammad Ali Akhaee, 2023 9th International Conference on Web Research (ICWR), 2023 [Link] [PDF]

Farsi CAPTCHA Recognition Using Attention-Based Convolutional Neural Network [Read More]

Matine Hajyan, Alireza Hosseni, Ramin Toosi, Mohammad Ali Akhaee, 2023 9th International Conference on Web Research (ICWR), 2023 [Link]

Face manifold: manifold learning for synthetic face generation [Read More]

Kimia Dinashi, Ramin Toosi, Mohammad Ali Akhaee, Multimedia Tools and Applications, 2023 [Link] [Code]

Hybrid Retrieval-Augmented Generation Approach for LLMs Query Response Enhancement [Read More]

Pouria Omrani, Alireza Hosseini, Kiana Hooshanfar, Zahra Ebrahimian, Ramin Toosi, Mohammad Ali Akhaee, 2024 10th International Conference on Web Research (ICWR), 2024 [Link]

Gender Recognition Based on Hand Images Employing Local and Global Shape Information [Read More]

Kiavash Jamshidi, Ramin Toosi, Mohammad Ali Akhaee, Computer, 2024 [Link]

CNN Autoencoder Resizer: A Power-Efficient LoS/NLoS Detector in MIMO-Enabled UAV Networks [Read More]

Azim Akhtarshenas, Navid Ayoobi, David Lopez-Perez, Ramin Toosi, Matin Amoozadeh, 2024 IEEE 35th International Symposium on Personal, Indoor and Mobile Radio Communications (PIMRC), 2024 [Link]

Knowledge Graph Based Retrieval-Augmented Generation for Multi-Hop Question Answering Enhancement [Read More]

Mahdi Amiri Shavaki, Pouria Omrani, Ramin Toosi, Mohammad Ali Akhaee, 2024 15th International Conference on Information and Knowledge Technology (IKT), 2024 [Link] [Code]

A Multi-Task Framework Using Mamba for Identity, Age, and Gender Classification from Hand Images [Read More]

Amirabbas Rezasoltani, Alireza Hosseini, Ramin Toosi, Mohammad Ali Akhaee, 2024 15th International Conference on Information and Knowledge Technology (IKT), 2024 [Link]

Microsaccade Selectivity as Discriminative Feature for Object Decoding [Read More]

Salar Nouri, Amirali Soltnai Tehrani, Niloufar Faridani, Ramin Toosi, Mohammad-Reza A Dehaqani, iScience, 2025 [Link] [PDF]

Few Shot Comic Character Re-identification [Read More]

Mahdi Kanani, Ramin Toosi, Under Review, 2025 [Link]

Judge a Book by its Cover: A Multimodal Approach to Book Genre Prediction [Read More]

Reza Toosi, Alireza Hosseini, Ramin Toosi, Mohammad Ali Akhaee, 2025 11th International Conference on Web Research (ICWR), 2025 [Link]

Brand Visibility in Packaging: A Deep Learning Approach for Logo Detection, Saliency-Map Prediction, and Logo Placement Analysis [Read More]

Alireza Hosseini, Kiana Hooshanfar, Pouria Omrani, Reza Toosi, Ramin Toosi, Zahra Ebrahimian, Mohammad Ali Akhaee, Discover Applied Sciences, 2025 [Link] [Code]

Unlocking Book Genre from Covers: A Multimodal Approach to Book Genre Prediction [Read More]

Reza Toosi, Alireza Hosseini, Ramin Toosi, Mohammad Ali Akhaee, International Journal of Web Research, 2025 [Link] [PDF]

Efficient Malicious UAV Detection Using Autoencoder-TSMamba Integration [Read More]

Azim Akhtarshenas, Ramin Toosi, David López-Pérez, Tohid Alizadeh, Alireza Hosseini, Iberian Conference on Pattern Recognition and Image Analysis, 2025 [Link]

The Spatial Frequency Representation Predicts Category Coding in the Inferior Temporal Cortex [Read More]

Ramin Toosi, Behnam Karami, Roxana Koushki, Farideh Shakerian, Jalaledin Noroozi, Ehsan Rezayat, Abdol-Hossein Vahabie, Mohammad Ali Akhaee, Mohammad-Reza A. Dehaqani, eLife, 2025 [Link]

A Comprehensive Mathematical and Applied Survey of Reinforcement Learning Algorithms [Read More]

Azim Akhtarshenas, Seyyed Hossein Mostafavi Tehrani, Ramin Toosi, Tohid Alizadeh, Mario Rico-Ibañez, Under Review, 2026 [Link]

Deep Gesture Recognition under Data Loss [Read More]

Amin Kajbaf, Ehsan Yazdian, Mohammad Ali Akhaee, Ramin Toosi, Saeed Gazor, IEEE Sensors Journal, 2026 [Link] [Code]

End-to-End Motion-Robust Gesture Recognition from Raw FMCW Data [Read More]

Amin Kajbaf, Ehsan Yazdian, Mohammad Ali Akhaee, Ramin Toosi, Saeed Gazor, Under Review, 2026 [Link] [Code]

Deep Multi-Task Locomotion-Invariant Gesture Recognition with FMCW Radar [Read More]

Amin Kajbaf, Ehsan Yazdian, Mohammad Ali Akhaee, Ramin Toosi, Saeed Gazor, Under Review, 2026 [Link] [Code]

Ramin Toosi

Sitemap

Pages

Page Not Found [Read More]

About Me [Read More]

Archive Layout with Content [Read More]

Blog Posts [Read More]

Posts by Category [Read More]

Posts by Collection [Read More]

Contact Me [Read More]

CV [Read More]

Education [Read More]

[Read More]

Markdown [Read More]

Page not in menu [Read More]

Page Archive [Read More]

Portfolio [Read More]

Projects [Read More]

Publications [Read More]

Sitemap [Read More]

Skills [Read More]

Posts by Tags [Read More]

Talk map [Read More]

Teaching [Read More]

Terms and Privacy Policy [Read More]

[Read More]

[Read More]

[Read More]

[Read More]

[Read More]

[Read More]

[Read More]

[Read More]

[Read More]

[Read More]

[Read More]

Jupyter notebook markdown generator [Read More]

[Read More]

[Read More]

Posts

Fine-Tuning LLaMA 3.2-11B-Vision for Product Descriptions

Diffusion From Scratch To Generate Cute Anime Faces!

Running Llama 3.2 in Rust

Boosting Model Efficiency with Quantization and Pruning in PyTorch

Let’s train a CNN in Rust with libtorch

Implementing K-Means in Rust: A 7x Speed Boost Over Python

Gender Bias in Multimodal Embeddings: An Example of OpenAI CLIP

Weak Supervision with Snorkel: Image Classification Example

Optimizing Python Code with LRU Cache: A Fibonacci Sequence Example

A += Assignment Puzzler in Python

Python Trick: Merging Two Lists with the zip Function

Lazy Predict: Explore lots of machine learning models at once

Accelerating Python code with Numba

What is Artificial Intelligence?

education

projects

Real-Time Face Search in Video [Read More]

Human Pose Estimation and Motion Capture [Read More]

Movie Genre Detection [Read More]

YouTube Downloader [Read More]

Twitter Job Title Prediction [Read More]

Video Content Creation and Curation [Read More]

Robust Offline Spike Sorter (ROSS) [Read More]

Mousi — Mouse Tracking and Behavioural Analysis [Read More]

Cage Camera for Rodent Behavioural Monitoring [Read More]

Brain-Inspired Spatial-Frequency-Aware Networks [Read More]

Weak Supervision with Snorkel [Read More]

K-Means Clustering in Rust [Read More]

Retrieval-Augmented Generation for LLM Systems [Read More]

CNN Training & Inference in Rust [Read More]

Model Optimization: Quantization and Pruning [Read More]

Llama Chatbot in Rust (llama.cpp) [Read More]

Diffusion Models from Scratch [Read More]

Deep Learning for MRI Reconstruction [Read More]

Spatial Frequency Representation in the Inferior Temporal Cortex [Read More]

Fine-Tuning LLaMA 3.2-Vision for Product Descriptions [Read More]

Multimodal Book Genre Prediction [Read More]

Persian Speech Recognition with Whisper [Read More]

Regex Engine & grep-Style Search (Rust) [Read More]

publications