如何在TensorFlow（不含Keras）中便捷设置预训练权重？

Got it, I totally get your frustration—Keras makes loading pretrained weights a breeze with model.set_weights() or layer.set_weights(), but when you're working with raw TensorFlow (no Keras layers involved), it feels like you have to get your hands a bit dirtier. Let's walk through the straightforward ways to do this, since it's all about working with TensorFlow's core tf.Variable objects under the hood.

1. Using tf.Module (TensorFlow's Native Modular Approach) #

tf.Module is TensorFlow's built-in way to organize your layers and variables, similar to Keras Model but more lightweight. Each layer's weights are just tf.Variable attributes, so you can directly assign your pretrained weights to them.

Here's a concrete example with a custom convolution layer:

import tensorflow as tf
import numpy as np

# Define a custom conv layer inheriting from tf.Module
class CustomConv(tf.Module):
    def __init__(self, num_filters, kernel_size):
        super().__init__()
        # Initialize placeholder weights (we'll overwrite these)
        self.kernel = tf.Variable(
            tf.random.normal([kernel_size, kernel_size, 3, num_filters]),
            dtype=tf.float32
        )
        self.bias = tf.Variable(
            tf.zeros([num_filters]),
            dtype=tf.float32
        )

    def __call__(self, inputs):
        return tf.nn.conv2d(inputs, self.kernel, strides=[1,1,1,1], padding="SAME") + self.bias

# Instantiate the layer
my_conv = CustomConv(num_filters=64, kernel_size=3)

# Load your pretrained weights (e.g., from numpy files)
pretrained_kernel = np.load("pretrained_conv_kernel.npy").astype(np.float32)
pretrained_bias = np.load("pretrained_conv_bias.npy").astype(np.float32)

# Assign the pretrained weights to the layer's variables
my_conv.kernel.assign(pretrained_kernel)
my_conv.bias.assign(pretrained_bias)

2. Directly Working with tf.Variable (No Wrappers) #

If you're not using tf.Module and just have standalone variables, the process is even simpler—you just call assign() directly on each variable.

# Define standalone weight variables
conv_kernel = tf.Variable(tf.random.normal([3,3,3,64]), dtype=tf.float32)
conv_bias = tf.Variable(tf.zeros([64]), dtype=tf.float32)

# Load pretrained weights
pretrained_kernel = np.load("pretrained_kernel.npy").astype(np.float32)
pretrained_bias = np.load("pretrained_bias.npy").astype(np.float32)

# Assign weights
conv_kernel.assign(pretrained_kernel)
conv_bias.assign(pretrained_bias)

Key Notes to Avoid Headaches #

• Shape & Type Matching: Make sure your pretrained weights (usually numpy arrays) have the exact same shape and data type as your TensorFlow variables. If your variable is float32, your numpy array must also be float32—mismatches will throw errors.
• Batch Assignment for Complex Models: For larger models, organize your variables into a dictionary (mapping variable names to tf.Variable objects) and your pretrained weights into a matching dictionary. Then you can loop through them to assign weights in bulk:

  # Dictionary of model variables
  model_vars = {
      "conv1/kernel": my_conv.kernel,
      "conv1/bias": my_conv.bias,
      # Add other variables here...
  }
  
  # Load pretrained weights as a dictionary (e.g., saved with numpy)
  pretrained_weights = np.load("pretrained_weights_dict.npy", allow_pickle=True).item()
  
  # Bulk assign
  for var_name, var in model_vars.items():
      var.assign(pretrained_weights[var_name])
  

Why Keras Feels Easier #

Keras layers just abstract away this manual variable targeting—layer.set_weights() takes a list of arrays that correspond to the layer's internal variables in a predefined order. Under the hood, it's doing the exact same assign() operations we're doing here!

内容的提问来源于stack exchange，提问作者Maverick Meerkat