name: "LeNet"

input: "data"
input_shape {
  dim: 1
  dim: 1
  dim: 28
  dim: 28
}
layer {
  name: "conv1"
  type: "Convolution"
  bottom: "data"
  top: "conv1"
  param {
    lr_mult: 1
  }
  param {
    lr_mult: 2
  }
  convolution_param {
    num_output: 20
    kernel_size: 5
    stride: 1
    weight_filler {
      type: "xavier"
    }
    bias_filler {
      type: "constant"
    }
  }
}
layer {
  name: "pool1"
  type: "Pooling"
  bottom: "conv1"
  top: "pool1"
  pooling_param {
    pool: MAX
    kernel_size: 2
    stride: 2
  }
}
layer {
  name: "conv2"
  type: "Convolution"
  bottom: "pool1"
  top: "conv2"
  param {
    lr_mult: 1
  }
  param {
    lr_mult: 2
  }
  convolution_param {
    num_output: 50
    kernel_size: 5
    stride: 1
    weight_filler {
      type: "xavier"
    }
    bias_filler {
      type: "constant"
    }
  }
}
layer {
  name: "pool2"
  type: "Pooling"
  bottom: "conv2"
  top: "pool2"
  pooling_param {
    pool: MAX
    kernel_size: 2
    stride: 2
  }
}
layer {
  name: "ip1"
  type: "InnerProduct"
  bottom: "pool2"
  top: "ip1"
  param {
    lr_mult: 1
  }
  param {
    lr_mult: 2
  }
  inner_product_param {
    num_output: 500
    weight_filler {
      type: "xavier"
    }
    bias_filler {
      type: "constant"
    }
  }
}
layer {
  name: "relu1"
  type: "ReLU"
  bottom: "ip1"
  top: "ip1"
}
layer {
  name: "ip2"
  type: "InnerProduct"
  bottom: "ip1"
  top: "ip2"
  param {
    lr_mult: 1
  }
  param {
    lr_mult: 2
  }
  inner_product_param {
    num_output: 10
    weight_filler {
      type: "xavier"
    }
    bias_filler {
      type: "constant"
    }
  }
}

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

# ================================================================ #
#                        Train a Sample Model                      #
# ================================================================ #
# 1. create data
mnist = input_data.read_data_sets(‘../MNIST‘, one_hot=True)
with tf.variable_scope(‘Input‘):
    tf_x = tf.placeholder(dtype=tf.float32, shape=[None, 28 * 28], name=‘x‘)
    image = tf.reshape(tf_x, [-1, 28, 28, 1], name=‘image‘)
    tf_y = tf.placeholder(dtype=tf.float32, shape=[None, 10], name=‘y‘)
    is_training = tf.placeholder(dtype=tf.bool, shape=None)

# 2. define Network
with tf.variable_scope(‘Net‘):
    """
    "SAME" 类型的padding:
    out_height = ceil(in_height / strides[1]); ceil向上取整
    out_width = ceil(in_width / strides[2])

    "VALID"类型的padding:
    out_height = ceil((in_height - filter_height + 1) / striders[1])
    out_width = ceil((in_width - filter_width + 1) / striders[2]
    """
    conv1 = tf.layers.conv2d(inputs=image, filters=20, kernel_size=5,
                             strides=1, padding=‘valid‘, activation=None, name=‘conv1‘)  # 20x24x24
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=2, strides=2, name=‘pool1‘)  # 20x12x12
    conv2 = tf.layers.conv2d(pool1, 50, 5, 1, ‘valid‘, activation=None, name=‘conv2‘)  # 50x8x8
    pool2 = tf.layers.max_pooling2d(conv2, 2, 2, name=‘pool2‘)  # 50x4x4
    pool2_flat = tf.reshape(pool2, [-1, 4 * 4 * 50])
    fc1 = tf.layers.dense(pool2_flat, 500, tf.nn.relu, name=‘ip1‘)  # 500
    predict = tf.layers.dense(fc1, 10, name=‘ip2‘)  # 10

# 3. define loss & accuracy
with tf.name_scope(‘loss‘):
    loss = tf.losses.softmax_cross_entropy(onehot_labels=tf_y, logits=predict, label_smoothing=0.01)
with tf.name_scope(‘accuracy‘):
    # tf.metrics.accuracy() 返回 累计[上次的平均accuracy， 这次的平均accuracy]
    accuracy = tf.metrics.accuracy(labels=tf.argmax(tf_y, axis=1), predictions=tf.argmax(predict, axis=1))[1]

# 4. define optimizer
with tf.name_scope(‘train‘):
    optimizer_op = tf.train.AdamOptimizer(1e-4).minimize(loss)

# 5. initialize
init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())

# 6.train
saver = tf.train.Saver()
save_path = ‘./leNet_mnist.ckpt‘

with tf.Session() as sess:
    sess.run(init_op)
    for step in range(11000):
        """
        mnist.train.num_examples=55000
        11000*100/mnist.train.num_examples=20epochs
        """
        batch_x, batch_y = mnist.train.next_batch(100)
        _, ls = sess.run([optimizer_op, loss],
                         feed_dict={tf_x: batch_x, tf_y: batch_y, is_training: True})
        if step % 100 == 0:
            acc_test = sess.run(accuracy,
                                feed_dict={tf_x: mnist.test.images, tf_y: mnist.test.labels,
                                           is_training: False})
            print(‘Step: ‘, step, ‘ | train loss: {:.4f} | test accuracy: {:.3f}‘.format(ls, acc_test))
            sess.run(tf.local_variables_initializer())  # 不加上这句的话 accuracy 就是个累积平均值了
    saver.save(sess, save_path)

# 7.test
with tf.Session() as sess:
    sess.run(init_op)
    saver.restore(sess, save_path)
    acc_test = sess.run(accuracy, feed_dict={tf_x: mnist.test.images,
                                             tf_y: mnist.test.labels,
                                             is_training: False})
    print(‘test accuracy: {}‘.format(acc_test)) # test accuracy: 0.991100013256073

from tensorflow.summary import FileWriter
sess = tf.Session()
tf.train.import_meta_graph("leNet_mnist.ckpt.meta")
FileWriter("__tb", sess.graph)

# ================================================================ #
#                Convert ckpt to pb & Compute FLOPs                #
# ================================================================ #
from tensorflow.python.framework import graph_util


def stats_graph(graph):
    flops = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.float_operation())
    params = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.trainable_variables_parameter())
    print(‘FLOPs: {};    Trainable params: {}‘.format(flops.total_float_ops, params.total_parameters))


with tf.Graph().as_default() as graph:
    # 1. Create Graph
    image = tf.Variable(initial_value=tf.random_normal([1, 28, 28, 1]))
    conv1 = tf.layers.conv2d(inputs=image, filters=20, kernel_size=5,
                             strides=1, padding=‘valid‘, activation=None, name=‘conv1‘)  # 20x24x24
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=2, strides=2, name=‘pool1‘)  # 20x12x12
    conv2 = tf.layers.conv2d(pool1, 50, 5, 1, ‘valid‘, activation=None, name=‘conv2‘)  # 50x8x8
    pool2 = tf.layers.max_pooling2d(conv2, 2, 2, name=‘pool2‘)  # 50x4x4
    pool2_flat = tf.reshape(pool2, [-1, 4 * 4 * 50])
    fc1 = tf.layers.dense(pool2_flat, 500, tf.nn.relu, name=‘ip1‘)  # 500
    predict = tf.layers.dense(fc1, 10, name=‘ip2‘)  # 10
    print(‘stats before freezing‘)
    stats_graph(graph)

    # 2. Freeze Graph
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        output_graph = graph_util.convert_variables_to_constants(sess, graph.as_graph_def(), [‘ip2/BiasAdd‘])
        with tf.gfile.GFile(‘LeNet_mnist.pb‘, "wb") as f:
            f.write(output_graph.SerializeToString())


def load_pb(pb):
    with tf.gfile.GFile(pb, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name=‘‘)
        return graph


# 3. Load Frozen Graph
graph = load_pb(‘LeNet_mnist.pb‘)
print(‘stats after freezing‘)
stats_graph(graph)

import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from torchsummary import summary

# Device configuration
device = torch.device(‘cpu‘) #torch.device(‘cuda: 0‘ if torch.cuda.is_available() else ‘cup‘)
print(device, torch.__version__)

# Hyper parameters
num_epochs = 5
num_classes = 10
batch_size = 100
learning_rate = 0.01

# MINST DATASET
train_dataset = torchvision.datasets.MNIST(root=‘H:/Other_Datasets/‘,
                                           train=True,
                                           transform=transforms.ToTensor(),
                                           download=True)

test_dataset = torchvision.datasets.MNIST(root=‘H:/Other_Datasets/‘,
                                          train=False,
                                          transform=transforms.ToTensor())

# Data loader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)


class LeNet(nn.Module):
    def __init__(self, in_channels, num_classes):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, 20, kernel_size=5, stride=1)  # 20x24x24
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # 20x12x12
        self.conv2 = nn.Conv2d(20, 50, kernel_size=5, stride=1)  # 50x8x8
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)  # 50x4x4
        self.fc1 = nn.Linear(50 * 4 * 4, 500)  # 500
        self.fc2 = nn.Linear(500, num_classes)  # 10

    def forward(self, input):
        out = self.conv1(input)
        out = self.pool1(out)
        out = self.conv2(out)
        out = self.pool2(out)
        out = out.reshape(out.size(0), -1)  # pytorch folow NCHW convention
        out = F.relu(self.fc1(out))
        out = self.fc2(out)
        return out


model = LeNet(1, num_classes).to(device)