Initializing tensorflow Variable with an array larger than 2GB

TensorFlow

Variable Initialization

Large Array

Deep Learning

Memory Management

Initializing tensorflow Variable with an array larger than 2GB

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Deep Learning TensorFlow & Keras Data Structures Complexity & Performance

Introduction

TensorFlow's protocol buffer serialization has a hard 2GB size limit per tensor. When you try to create a tf.Variable or tf.constant from a NumPy array larger than 2GB, TensorFlow raises ValueError: Cannot create a tensor proto whose content is larger than 2GB. The fix is to use tf.Variable with an initializer function, load the data in chunks via variable.assign(), use tf.data.Dataset, or store the data in a file and read it at runtime.

The Problem

python

1import numpy as np
2import tensorflow as tf
3
4# A 2.5GB float32 array
5large_array = np.random.randn(250_000_000, 3).astype(np.float32)
6print(f"Array size: {large_array.nbytes / 1e9:.1f} GB")  # 3.0 GB
7
8# This fails
9var = tf.Variable(large_array)
10# ValueError: Cannot create a tensor proto whose content is larger than 2GB

The limit comes from Protocol Buffers (protobuf), which TensorFlow uses internally for tensor serialization. Protobuf messages cannot exceed 2GB.

Fix 1: Use an Initializer Function

Instead of passing the array directly, pass a callable that returns the data:

python

1import numpy as np
2import tensorflow as tf
3
4large_array = np.random.randn(250_000_000, 3).astype(np.float32)
5
6# Use a lambda initializer — avoids protobuf serialization
7var = tf.Variable(
8    initial_value=lambda: tf.constant(large_array[:125_000_000, :]),
9    trainable=False,
10)
11
12# Or split into chunks and concatenate
13def init_fn():
14    chunk_size = 50_000_000
15    chunks = []
16    for i in range(0, len(large_array), chunk_size):
17        chunks.append(tf.constant(large_array[i:i + chunk_size]))
18    return tf.concat(chunks, axis=0)
19
20var = tf.Variable(initial_value=init_fn, trainable=False)

Fix 2: Assign in Chunks

Create the variable with a placeholder shape, then assign the data in pieces:

python

1import numpy as np
2import tensorflow as tf
3
4large_array = np.random.randn(250_000_000, 3).astype(np.float32)
5
6# Create variable with zeros, then assign chunks
7var = tf.Variable(tf.zeros((250_000_000, 3), dtype=tf.float32), trainable=False)
8
9chunk_size = 50_000_000
10for i in range(0, len(large_array), chunk_size):
11    end = min(i + chunk_size, len(large_array))
12    var[i:end].assign(tf.constant(large_array[i:end]))
13    print(f"Assigned rows {i:,} to {end:,}")

Fix 3: Use tf.data.Dataset

For training data that does not need to be a single variable, load it via tf.data:

python

1import numpy as np
2import tensorflow as tf
3
4large_array = np.random.randn(250_000_000, 3).astype(np.float32)
5labels = np.random.randint(0, 10, size=(250_000_000,))
6
7# Create dataset from generator — no size limit
8def data_generator():
9    for i in range(len(large_array)):
10        yield large_array[i], labels[i]
11
12dataset = tf.data.Dataset.from_generator(
13    data_generator,
14    output_signature=(
15        tf.TensorSpec(shape=(3,), dtype=tf.float32),
16        tf.TensorSpec(shape=(), dtype=tf.int64),
17    ),
18)
19
20# Batch and prefetch for training
21dataset = dataset.batch(1024).prefetch(tf.data.AUTOTUNE)

Fix 4: Load from File at Runtime

Store large arrays in NumPy, HDF5, or TFRecord files:

python

1import numpy as np
2import tensorflow as tf
3
4# Save to disk
5large_array = np.random.randn(250_000_000, 3).astype(np.float32)
6np.save("large_embedding.npy", large_array)
7
8# Load at runtime using memory mapping — no 2GB limit
9def load_embedding():
10    data = np.load("large_embedding.npy", mmap_mode="r")
11    return tf.constant(data)  # Still limited per chunk
12
13# Better: load in chunks
14def load_embedding_chunked(path, shape, dtype=tf.float32):
15    data = np.load(path, mmap_mode="r")
16    var = tf.Variable(tf.zeros(shape, dtype=dtype), trainable=False)
17    chunk_size = 50_000_000
18    for i in range(0, shape[0], chunk_size):
19        end = min(i + chunk_size, shape[0])
20        var[i:end].assign(tf.constant(data[i:end]))
21    return var
22
23embedding = load_embedding_chunked("large_embedding.npy", (250_000_000, 3))

Fix 5: TF1 — Use tf.placeholder with feed_dict

In TensorFlow 1.x, use placeholders and feed the data:

python

1import numpy as np
2import tensorflow as tf
3
4large_array = np.random.randn(250_000_000, 3).astype(np.float32)
5
6# TF1 approach
7placeholder = tf.placeholder(tf.float32, shape=[None, 3])
8var = tf.Variable(tf.zeros([250_000_000, 3]), trainable=False)
9assign_op = var.assign(placeholder)
10
11with tf.Session() as sess:
12    sess.run(tf.global_variables_initializer())
13    # Feed in chunks
14    chunk_size = 50_000_000
15    for i in range(0, len(large_array), chunk_size):
16        chunk = large_array[i:i + chunk_size]
17        sess.run(var[i:i + chunk_size].assign(chunk))

Large Embedding Tables

The most common use case for >2GB variables is embedding tables in recommendation models:

python

1import tensorflow as tf
2
3# 100M embeddings × 64 dimensions = ~25GB in float32
4vocab_size = 100_000_000
5embedding_dim = 64
6
7# Use tf.keras.layers.Embedding with a file-backed initializer
8class FileInitializer(tf.keras.initializers.Initializer):
9    def __init__(self, path, chunk_size=10_000_000):
10        self.path = path
11        self.chunk_size = chunk_size
12
13    def __call__(self, shape, dtype=None):
14        data = np.load(self.path, mmap_mode="r")
15        result = tf.Variable(tf.zeros(shape, dtype=dtype or tf.float32))
16        for i in range(0, shape[0], self.chunk_size):
17            end = min(i + self.chunk_size, shape[0])
18            result[i:end].assign(tf.constant(data[i:end]))
19        return result
20
21embedding_layer = tf.keras.layers.Embedding(
22    vocab_size,
23    embedding_dim,
24    embeddings_initializer=FileInitializer("embeddings.npy"),
25)

Common Pitfalls

Trying to serialize the entire array at once: tf.constant(large_array) serializes the array through protobuf, hitting the 2GB limit. Always chunk arrays larger than ~500M elements (2GB / 4 bytes per float32) when creating tensors.
Forgetting memory-mapped loading: np.load("file.npy") loads the entire file into RAM. For very large files, use np.load("file.npy", mmap_mode="r") to memory-map it, then load chunks into TensorFlow incrementally.
Creating tf.constant in a loop without clearing: Each tf.constant(chunk) creates a graph node in TF1 or a tensor in TF2. Creating thousands of constants in a loop wastes memory. Assign directly to variable slices instead of accumulating constants.
Not using float16 to halve the size: If full float32 precision is not needed, converting to float16 halves the memory and may bring the tensor under 2GB. Use large_array.astype(np.float16) and tf.float16.
Ignoring the 2GB limit in SavedModel: Even if you initialize a variable with chunks, saving with tf.saved_model.save() serializes each variable into protobuf, hitting the limit again. Use tf.train.Checkpoint instead, which stores variables in a chunked format without the protobuf size limit.

Summary

TensorFlow cannot create a single tensor proto larger than 2GB (protobuf limitation)
Initialize large variables with a callable (lambda) or assign data in chunks via var[i:j].assign()
Use tf.data.Dataset with generators for large training data that does not need to be a single tensor
Store large arrays in .npy files and load with memory mapping (mmap_mode="r")
Use tf.train.Checkpoint instead of SavedModel to save models with large variables
Consider float16 to halve memory usage and potentially stay under the 2GB limit

DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Initializing tensorflow Variable with an array larger than 2GB

ML System Design practice on Codemia

Introduction

The Problem

Fix 1: Use an Initializer Function

Fix 2: Assign in Chunks

Fix 3: Use tf.data.Dataset

Fix 4: Load from File at Runtime

Fix 5: TF1 — Use tf.placeholder with feed_dict

Large Embedding Tables

Common Pitfalls

Summary

Related reading

DSA Fundamentals

Track what you have practised

ML System Design practice on Codemia