transfer learning

Author: MorvanZhou

README.md

@@ -36,7 +36,7 @@ All methods mentioned below have their video and text tutorial in Chinese. Visit
   * [AutoEncoder](https://github.com/MorvanZhou/Tensorflow-Tutorial/blob/master/tutorial-contents/404_AutoEncoder.py)
   * [DQN Reinforcement Learning](https://github.com/MorvanZhou/Tensorflow-Tutorial/blob/master/tutorial-contents/405_DQN_reinforcement_learning.py)
   * [GAN (Generative Adversarial Nets)](https://github.com/MorvanZhou/Tensorflow-Tutorial/blob/master/tutorial-contents/406_GAN.py) / [Conditional GAN](https://github.com/MorvanZhou/Tensorflow-Tutorial/blob/master/tutorial-contents/406_conditional_GAN.py)
-
+  * [Transfer Learning](https://github.com/MorvanZhou/Tensorflow-Tutorial/blob/master/tutorial-contents/407_transfer_learning.py)
 * Others (WIP)
   * [Dropout](https://github.com/MorvanZhou/Tensorflow-Tutorial/blob/master/tutorial-contents/501_dropout.py)
   * [Batch Normalization](https://github.com/MorvanZhou/Tensorflow-Tutorial/blob/master/tutorial-contents/502_batch_normalization.py)

tutorial-contents/407_transfer_learning.py

@@ -0,0 +1,195 @@
+"""
+This is a simple example of transfer learning using VGG.
+Fine tune a CNN from a classifier to regressor.
+Generate some fake data for describing cat and tiger length.
+
+Fake length setting:
+Cat - Normal distribution (40, 8)
+Tiger - Normal distribution (100, 30)
+
+The VGG model and parameters are adopted from:
+https://github.com/machrisaa/tensorflow-vgg
+
+Learn more, visit my tutorial site: [莫烦Python](https://morvanzhou.github.io)
+"""
+
+from urllib.request import urlretrieve
+import os
+import numpy as np
+import tensorflow as tf
+import skimage.io
+import skimage.transform
+import matplotlib.pyplot as plt
+
+
+def download():     # download tiger and kittycat image
+    categories = ['tiger', 'kittycat']
+    for category in categories:
+        os.makedirs('./for_transfer_learning/data/%s' % category, exist_ok=True)
+        with open('./for_transfer_learning/imagenet_%s.txt' % category, 'r') as file:
+            urls = file.readlines()
+            n_urls = len(urls)
+            for i, url in enumerate(urls):
+                try:
+                    urlretrieve(url.strip(), './data/%s/%s' % (category, url.strip().split('/')[-1]))
+                    print('%s %i/%i' % (category, i, n_urls))
+                except:
+                    print('%s %i/%i' % (category, i, n_urls), 'no image')
+
+
+def load_img(path):
+    img = skimage.io.imread(path)
+    img = img / 255.0
+    # print "Original Image Shape: ", img.shape
+    # we crop image from center
+    short_edge = min(img.shape[:2])
+    yy = int((img.shape[0] - short_edge) / 2)
+    xx = int((img.shape[1] - short_edge) / 2)
+    crop_img = img[yy: yy + short_edge, xx: xx + short_edge]
+    # resize to 224, 224
+    resized_img = skimage.transform.resize(crop_img, (224, 224))[None, :, :, :]   # shape [1, 224, 224, 3]
+    return resized_img
+
+
+def load_data():
+    imgs = {'tiger': [], 'kittycat': []}
+    for k in imgs.keys():
+        dir = './data/' + k
+        for file in os.listdir(dir):
+            if not file.lower().endswith('.jpg'):
+                continue
+            try:
+                resized_img = load_img(os.path.join(dir, file))
+            except OSError:
+                continue
+            imgs[k].append(resized_img)    # [1, height, width, depth] * n
+            if len(imgs[k]) == 400:        # only use 400 imgs to reduce my memory load
+                break
+    # fake length data for tiger and cat
+    tigers_y = np.maximum(20, np.random.randn(len(imgs['tiger']), 1) * 30 + 100)
+    cat_y = np.maximum(10, np.random.randn(len(imgs['kittycat']), 1) * 8 + 40)
+    return imgs['tiger'], imgs['kittycat'], tigers_y, cat_y
+
+
+class Vgg16:
+    vgg_mean = [103.939, 116.779, 123.68]
+
+    def __init__(self, vgg16_npy_path=None, restore_from=None):
+        # pre-trained parameters
+        try:
+            self.data_dict = np.load(vgg16_npy_path, encoding='latin1').item()
+        except FileNotFoundError:
+            print('Please download VGG16 parameters at here https://mega.nz/#!YU1FWJrA!O1ywiCS2IiOlUCtCpI6HTJOMrneN-Qdv3ywQP5poecM')
+
+        self.tfx = tf.placeholder(tf.float32, [None, 224, 224, 3])
+        self.tfy = tf.placeholder(tf.float32, [None, 1])
+
+        # Convert RGB to BGR
+        red, green, blue = tf.split(axis=3, num_or_size_splits=3, value=self.tfx * 255.0)
+        bgr = tf.concat(axis=3, values=[
+            blue - self.vgg_mean[0],
+            green - self.vgg_mean[1],
+            red - self.vgg_mean[2],
+        ])
+
+        # pre-trained VGG layers are fixed in fine-tune
+        conv1_1 = self.conv_layer(bgr, "conv1_1")
+        conv1_2 = self.conv_layer(conv1_1, "conv1_2")
+        pool1 = self.max_pool(conv1_2, 'pool1')
+
+        conv2_1 = self.conv_layer(pool1, "conv2_1")
+        conv2_2 = self.conv_layer(conv2_1, "conv2_2")
+        pool2 = self.max_pool(conv2_2, 'pool2')
+
+        conv3_1 = self.conv_layer(pool2, "conv3_1")
+        conv3_2 = self.conv_layer(conv3_1, "conv3_2")
+        conv3_3 = self.conv_layer(conv3_2, "conv3_3")
+        pool3 = self.max_pool(conv3_3, 'pool3')
+
+        conv4_1 = self.conv_layer(pool3, "conv4_1")
+        conv4_2 = self.conv_layer(conv4_1, "conv4_2")
+        conv4_3 = self.conv_layer(conv4_2, "conv4_3")
+        pool4 = self.max_pool(conv4_3, 'pool4')
+
+        conv5_1 = self.conv_layer(pool4, "conv5_1")
+        conv5_2 = self.conv_layer(conv5_1, "conv5_2")
+        conv5_3 = self.conv_layer(conv5_2, "conv5_3")
+        pool5 = self.max_pool(conv5_3, 'pool5')
+
+        # detach original VGG fc layers and
+        # reconstruct your own fc layers serve for your own purpose
+        self.flatten = tf.reshape(pool5, [-1, 7*7*512])
+        self.fc6 = tf.layers.dense(self.flatten, 256, tf.nn.relu, name='fc6')
+        self.out = tf.layers.dense(self.fc6, 1, name='out')
+
+        self.sess = tf.Session()
+        if restore_from:
+            saver = tf.train.Saver()
+            saver.restore(self.sess, restore_from)
+        else:   # training graph
+            self.loss = tf.losses.mean_squared_error(labels=self.tfy, predictions=self.out)
+            self.train_op = tf.train.RMSPropOptimizer(0.001).minimize(self.loss)
+            self.sess.run(tf.global_variables_initializer())
+
+    def max_pool(self, bottom, name):
+        return tf.nn.max_pool(bottom, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name=name)
+
+    def conv_layer(self, bottom, name):
+        with tf.variable_scope(name):   # filter is constant
+            conv = tf.nn.conv2d(bottom, self.data_dict[name][0], [1, 1, 1, 1], padding='SAME')
+            lout = tf.nn.relu(tf.nn.bias_add(conv, self.data_dict[name][1]))
+            return lout
+
+    def train(self, x, y):
+        loss, _ = self.sess.run([self.loss, self.train_op], {self.tfx: x, self.tfy: y})
+        return loss
+
+    def predict(self, paths):
+        fig, axs = plt.subplots(1, 2)
+        for i, path in enumerate(paths):
+            x = load_img(path)
+            length = self.sess.run(self.out, {self.tfx: x})
+            axs[i].imshow(x[0])
+            axs[i].set_title('Len: %.1f cm' % length)
+            axs[i].set_xticks(()); axs[i].set_yticks(())
+        plt.show()
+
+    def save(self, path='./for_transfer_learning/model/transfer_learn'):
+        saver = tf.train.Saver()
+        saver.save(self.sess, path, write_meta_graph=False)
+
+
+def train():
+    tigers_x, cats_x, tigers_y, cats_y = load_data()
+
+    # plot fake length distribution
+    plt.hist(tigers_y, bins=20, label='Tigers')
+    plt.hist(cats_y, bins=10, label='Cats')
+    plt.legend()
+    plt.xlabel('length')
+    plt.show()
+
+    xs = np.concatenate(tigers_x + cats_x, axis=0)
+    ys = np.concatenate((tigers_y, cats_y), axis=0)
+
+    vgg = Vgg16(vgg16_npy_path='./for_transfer_learning/vgg16.npy')
+    print('Net built')
+    for i in range(100):
+        b_idx = np.random.randint(0, len(xs), 6)
+        train_loss = vgg.train(xs[b_idx], ys[b_idx])
+        print(i, 'train loss: ', train_loss)
+
+    vgg.save('./for_transfer_learning/model/transfer_learn')      # save learned fc layers
+
+
+def eval():
+    vgg = Vgg16(vgg16_npy_path='./for_transfer_learning/vgg16.npy',
+                restore_from='./for_transfer_learning/model/transfer_learn')
+    vgg.predict(
+        ['./for_transfer_learning/data/kittycat/000129037.jpg', './for_transfer_learning/data/tiger/391412.jpg'])
+
+
+if __name__ == '__main__':
+    # download()
+    # train()
+    eval()