Skip to main content

Chapter 12.3 - Mapping TF to PyTorch

Overview

Sometimes, an incredibly smart AI was trained by Google in a language called TensorFlow. But we are building our model in PyTorch! We need to translate the AI's brain (its weights) from TensorFlow language to PyTorch language.

🎯 Why we do it​

Rationale

OpenAI spent millions of dollars training GPT-2 in TensorFlow. We don't want to spend millions of dollars training our own from scratch! So, we just download their smart brain and carefully map their TensorFlow names to our PyTorch names.

🛠️ How we do it​

Methodology

We create a dictionary (like a translation book). We say, "Hey, whenever you see a weight named h0/attn/c_attn/w in TensorFlow, put it inside our PyTorch layer named transformer.h.0.attn.c_attn.weight."

# A simple translation dictionary!
name_map = {
    "h0/attn/c_attn/w": "transformer.h.0.attn.c_attn.weight",
    "h0/ln_1/b": "transformer.h.0.ln_1.bias"
}

print("We know exactly where to plug in the smart weights!")