tensorflow 加載預訓練模型進行 finetune 的操作解析

這是一篇需要仔細思考的博客；

 

預訓練模型

tensorflow 在 1.0 之后移除了 models 模塊，這個模塊實現了很多模型，並提供了部分預訓練模型的權重；

圖像識別模型的權重下載地址  https://github.com/tensorflow/models/tree/master/research/slim

 

模型加載

首先需要了解模型保存的形式，包含了 checkpoint、data、meta 等文件；

模型加載不僅可以從 data 加載訓練好的權重，還可以從 meta 加載計算圖，

加載計算圖我們可以理解為引入了 計算節點和變量，引入變量很重要，這樣我們無需自己去創造變量，

加載計算圖返回的是個 Saver 對象，如果沒有通過 加載圖引入變量，也沒有自己創造變量，是無法創建 Saver 對象的，沒有 Saver 對象，就無法加載預訓練權重；

 

下面我用代碼解釋上面的邏輯；

首先做個數據准備：寫個最簡單的計算圖，然后保存

with tf.name_scope('scope1'):
    w1 = tf.Variable(1, name='w1')
    v1 = tf.Variable(2, name='v1')

with tf.name_scope('scope2'):
    w2 = tf.Variable(3, name='w2')
    v2 = tf.Variable(4, name='v2')

out = tf.add(w1*v1, w2*v2)
init = tf.global_variables_initializer()

saver = tf.train.Saver()

sess = tf.Session()
sess.run(init)
print(sess.run(out))
saver.save(sess, 'data/test.ckpt')

記住這里 w1 的值為 1，后面有用

 

加載預訓練權重有兩種思路

1. 先加載計算圖，再加載權重

2. 自己創建計算圖，再加載權重

總結一句話就是 先有變量(創建圖就必須創建變量)，然后創建 Saver 對象，通過 Saver 對象加載權重

 

先加載計算圖，再加載權重

加載計算圖用下面的方法

def import_meta_graph(meta_graph_or_file,
                      clear_devices=False,
                      import_scope=None,
                      **kwargs):
  """Recreates a Graph saved in a `MetaGraphDef` proto."""

加載計算圖后，通過 sess.graph 獲取圖，通過 get_tensor_by_name 等方法獲取保存的變量

 

簡單舉例

## 加載了圖中的各個節點，相當於引入變量
saver = tf.train.import_meta_graph('data/test.ckpt.meta')       # 從 meta 文件直接加載圖，返回一個存儲器
print(type(saver))      # <class 'tensorflow.python.training.saver.Saver'>

### 如果沒有變量，直接創建存儲器，會報錯的
# saver = tf.train.Saver()      ### 報錯 ValueError: No variables to save

sess = tf.Session()
saver.restore(sess, 'data/test.ckpt')   ### 獲取圖權重
# print(sess.run('w1'))           ### 直接這樣報錯

graph = sess.graph  ### 獲取加載的圖
print(sess.run(graph.get_tensor_by_name('scope1/w1:0')))        # 1     ### 獲取圖的 Tensor
# print(sess.run(graph.get_tensor_by_name('w1:0')))       ### 報錯 KeyError: "The name 'w1:0' refers to a Tensor which does not exist. The operation, 'w1', does not exist in the graph."
sess.close()

 

自己創建計算圖，再加載權重

注意，自己創建的計算圖要與保存的計算圖一致，包括 網絡結構、作用域、變量名等

##### 上面是直接加載圖，引入變量，從而創建存儲器
##### 這里我們不加載，自己創建一個圖，從而創建存儲器
### 注意，圖的結構要與 checkpoint 中的一致，作用域、變量名 都要一樣
with tf.name_scope('scope1'):
    w1 = tf.Variable(333, name='w1')
    v1 = tf.Variable(2, name='v1')

with tf.name_scope('scope2'):
    w2 = tf.Variable(3, name='w2')
    v2 = tf.Variable(4, name='v2')

init = tf.global_variables_initializer()

saver = tf.train.Saver()

sess = tf.Session()
sess.run(init)
saver.restore(sess, 'data/test.ckpt')

graph = tf.get_default_graph()      ### 獲取默認圖
print(sess.run(graph.get_tensor_by_name('scope1/w1:0')))        # 1

這里我們創建計算圖是把 w1 賦值 333，而加載的 w1 仍然是保存時的 1，說明加載成功了

 

加載預訓練權重進行 finetune

finetune 我會單獨寫一篇博客，這里不細講，它大致可以分兩種：

1. 加載部分權重，其他權重正常初始化，然后優化所有參數；

2. 加載部分權重，其他權重正常初始化，然后優化非加載的參數，相當於固定部分權重，優化另一部分權重；

需要加載的權重 肯定通過 預訓練模型獲取，而其他權重則既可以通過預訓練的模型獲取，也可以自己創建，這點在上一章已經講清楚了，下面我們為了方便，就直接加載計算圖了；

 

固定部分權重是個難點，它的思路有兩點：

1. 先加載這部分權重，然后把這部分權重經過前向計算，得到一個新的 Input_new，然后把這個 Input 作為后續網絡的輸入，反向傳播時傳到 Input 肯定就停止了；

2. 分別加載需要訓練的參數 train_var 和不需要訓練的參數 fixed_var，然后在 優化器的 optimizer.minimize 方法中指定 var_list 為 train_var

def minimize(self, loss, global_step=None, var_list=None,
               gate_gradients=GATE_OP, aggregation_method=None,
               colocate_gradients_with_ops=False, name=None,
               grad_loss=None):
    """Add operations to minimize `loss` by updating `var_list`.
      var_list: Optional list or tuple of `Variable` objects to update to
        minimize `loss`.  Defaults to the list of variables collected in
        the graph under the key `GraphKeys.TRAINABLE_VARIABLES`."""

 

創造計算圖，先加載 fixed_var，再添加新的網絡層

demo 如下：只是偽代碼哦

tf.reset_default_graph()        ### 這句暫時可忽略

# 構建計算圖
images = tf.placeholder(tf.float32,(None,224,224,3))
with tf.contrib.slim.arg_scope(mobilenet_v2.training_scope(is_training=False)):
    logits, endpoints = mobilenet_v2.mobilenet(images,depth_multiplier=1.4)

with tf.variable_scope("finetune_layers"):
    mobilenet_tensor = tf.get_default_graph().get_tensor_by_name("MobilenetV2/expanded_conv_14/output:0")       # 獲取目標張量，取出mobilenet中指定層的張量

    # 將張量作為新的 Input 向新層傳遞
    x = tf.layers.Conv2D(filters=256,kernel_size=3,name="conv2d_1")(mobilenet_tensor)
    x = tf.nn.relu(x,name="relu_1")
    x = tf.layers.Conv2D(filters=256,kernel_size=3,name="conv2d_2")(x)
    x = tf.layers.Conv2D(10,3,name="conv2d_3")(x)
    predictions = tf.reshape(x, (-1,10))

 

分別獲取 train_var 和 fixed_var，在 minimize 中指定 var_list

demo 如下：只是偽代碼哦

#### 不重要的函數
def get_var_list(target_tensor=None):
    '''獲取指定變量列表 var_list 的函數；
       具體怎么干的，無需關心，只需要知道它的作用是獲取一批權重
    '''
    if target_tensor==None:
        target_tensor = r"MobilenetV2/expanded_conv_14/output:0"
    target = target_tensor.split("/")[1]
    all_list = []
    all_var = []

    for var in tf.global_variables():
        if var != []:
            all_list.append(var.name)
            all_var.append(var)
    try:
        all_list = list(map(lambda x:x.split("/")[1],all_list))
        # 查找對應變量作用域的索引
        ind = all_list[::-1].index(target)
        ind = len(all_list) -  ind - 1
        print(ind)
        del all_list
        return all_var[:ind+1]
    except:
        print("target_tensor is not exist!")
        
#### 下面這一堆仔細看
x_train = np.random.random(size=(141,224,224,3))
y_train = to_categorical(label_fake,10)

y_label = tf.placeholder(tf.int32, (None,10))

### 收集變量作用域 finetune_layers 內的可訓練變量，作為 train_var
train_var = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,scope="finetune_layers")

loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=y_label,logits=logits)
### 定義優化方法，用 var_list 指定需要更新的權重，此時僅更新 train_var 權重
train_step = tf.train.GradientDescentOptimizer(0.1).minimize(loss,var_list=train_var)

epochs = 10
batch_size = 16

# 目標張量名稱，要獲取變量列表 fixed_var
target_tensor = "MobilenetV2/expanded_conv_14/output:0"
fixed_var = get_var_list(target_tensor)
saver = tf.train.Saver(var_list=fixed_var)

with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
    writer = tf.summary.FileWriter(r"./logs", sess.graph)
    ## 初始化 train_var， 使用初始化指定函數
    sess.run(tf.variables_initializer(var_list=train_var))
    saver.restore(sess,tf.train.latest_checkpoint("./model_ckpt/mobilenet_v2"))

    for i in range(2000):
        start = (i*batch_size) % x_train.shape[0]
        end = min(start+batch_size, x_train.shape[0])
        _, merge, losses = sess.run([train_step,merge_all,loss], feed_dict={images:x_train[start:end], y_label:y_train[start:end]})
        if i%100==0: writer.add_summary(merge, i)

重點就下面 3 句

train_step = tf.train.GradientDescentOptimizer(0.1).minimize(loss,var_list=train_var)
sess.run(tf.variables_initializer(var_list=train_var))    
saver.restore(sess,tf.train.latest_checkpoint("./model_ckpt/mobilenet_v2"))

深入理解一下：

1. 如果初始化了 train_var 后，又加載了所有預訓練變量，也就是說 train_var 的初始值也是 預訓練的，而不是平常的 全0、全1、高斯分布等，這是可以的；

2. 如果先 加載所有預訓練變量，然后初始化 train_var，也是可以的，因為 fixed_var 沒有重新初始化，還是 預訓練的值，而 train_var 初始值是多少沒那么重要；

3. 如果初始化 train_var，加載 fixed_var，則誰先誰后無所謂；

4. 如果是 先用 tf.global_variables_initializer() 初始化全部參數，再加載全部預訓練參數，也是可以的；

5. 如果先加載全部預訓練參數，在用 tf.global_variables_initializer() 初始化全部參數，是不可以的，因為 fixed_var 也被初始化了；

上述代碼采用的是第 3 種，指定了只加載 fixed_var

saver = tf.train.Saver(var_list=fixed_var)

 

 

 

參考資料：

https://zhuanlan.zhihu.com/p/42183653