pytorch_parser.py

import torch
import torch.nn as nn
from typing import List, Dict, Any
from ir.hir import HIR


class PyTorchParser:
    """
    PyTorch model parser that converts PyTorch models to HIR.
    """
    
    def parse(self, model: nn.Module, input_shape=(1,4)) -> HIR:
        """
        Parse a PyTorch model into HIR representation.
        
        Args:
            model: PyTorch nn.Module to parse
            input_shape: Input tensor shape
            
        Returns:
            HIR representation of the model
        """
        return parse_pytorch_model(model, input_shape)


def parse_pytorch_model(model: nn.Module, input_shape=(1,4)) -> HIR:
    """
    Parser: extracts named modules and creates HIR.
    Now supports Linear, ReLU, Conv2d, AdaptiveAvgPool2d, Flatten, BatchNorm2d.
    """
    inputs = [{"name": "input", "dtype": "float32", "shape": list(input_shape)}]
    outputs = ["output"]
    blocks = []
    nodes = []
    for name, module in model.named_modules():
        # skip the top-level container
        if name == "":
            continue
        if isinstance(module, nn.Linear):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "Linear",
                "name": name,
                "params": {"in": module.in_features, "out": module.out_features},
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.Conv2d):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "Conv2d",
                "name": name,
                "params": {
                    "in_channels": module.in_channels,
                    "out_channels": module.out_channels,
                    "kernel_size": list(module.kernel_size) if isinstance(module.kernel_size, tuple) else [module.kernel_size, module.kernel_size],
                    "stride": list(module.stride) if isinstance(module.stride, tuple) else [module.stride, module.stride],
                    "padding": list(module.padding) if isinstance(module.padding, tuple) else [module.padding, module.padding],
                    "groups": module.groups,
                    "bias": module.bias is not None
                },
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.ReLU):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "ReLU",
                "name": name,
                "params": {},
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.AdaptiveAvgPool2d):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "AdaptiveAvgPool2d",
                "name": name,
                "params": {"output_size": list(module.output_size)},
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.MaxPool2d):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "MaxPool2d",
                "name": name,
                "params": {
                    "kernel_size": list(module.kernel_size) if isinstance(module.kernel_size, tuple) else [module.kernel_size, module.kernel_size],
                    "stride": list(module.stride) if isinstance(module.stride, tuple) else [module.stride, module.stride],
                    "padding": list(module.padding) if isinstance(module.padding, tuple) else [module.padding, module.padding]
                },
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.AvgPool2d):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "AvgPool2d",
                "name": name,
                "params": {
                    "kernel_size": list(module.kernel_size) if isinstance(module.kernel_size, tuple) else [module.kernel_size, module.kernel_size],
                    "stride": list(module.stride) if isinstance(module.stride, tuple) else [module.stride, module.stride],
                    "padding": list(module.padding) if isinstance(module.padding, tuple) else [module.padding, module.padding]
                },
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.Dropout):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "Dropout",
                "name": name,
                "params": {"p": module.p},
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.Flatten):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "Flatten",
                "name": name,
                "params": {
                    "start_dim": module.start_dim,
                    "end_dim": module.end_dim
                },
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.BatchNorm2d):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "BatchNorm2d",
                "name": name,
                "params": {
                    "num_features": module.num_features,
                    "eps": module.eps,
                    "momentum": module.momentum
                },
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.MultiheadAttention):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "MultiheadAttention",
                "name": name,
                "params": {
                    "embed_dim": module.embed_dim,
                    "num_heads": module.num_heads,
                    "dropout": module.dropout,
                    "bias": module.in_proj_bias is not None,
                    "add_bias_kv": module.bias_k is not None,
                    "add_zero_attn": module.add_zero_attn,
                    "kdim": module.kdim,
                    "vdim": module.vdim
                },
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.LayerNorm):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "LayerNorm",
                "name": name,
                "params": {
                    "normalized_shape": list(module.normalized_shape),
                    "eps": module.eps,
                    "elementwise_affine": module.elementwise_affine
                },
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.GELU):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "GELU",
                "name": name,
                "params": {
                    "approximate": getattr(module, "approximate", "none")
                },
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.Embedding):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "Embedding",
                "name": name,
                "params": {
                    "num_embeddings": module.num_embeddings,
                    "embedding_dim": module.embedding_dim,
                    "padding_idx": module.padding_idx,
                    "max_norm": module.max_norm,
                    "norm_type": module.norm_type,
                    "scale_grad_by_freq": module.scale_grad_by_freq,
                    "sparse": module.sparse
                },
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.SiLU):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "SiLU",
                "name": name,
                "params": {},
                "inputs": [], "outputs": []
            })
        elif isinstance(module, nn.Mish):
            nodes.append({
                "id": name.replace(".", "_"),
                "op": "Mish",
                "name": name,
                "params": {},
                "inputs": [], "outputs": []
            })
        else:
            # ignore other modules for now
            continue

    blocks.append({"id": "main", "type": "sequential", "nodes": nodes, "metadata": {}})
    hir = HIR(version="itc-v0.1", kind="model", name=model.__class__.__name__,
              framework="pytorch", inputs=inputs, outputs=outputs, blocks=blocks,
              metadata={"note": "generated by pytorch_parser"})
    return hir