SRCNN 代码解读【model.py】
最近在学习SRCNN，阅读代码做好笔记
代码下载链接https://github.com/tegg89/SRCNN-Tensorflow
下面开始
from utils import (
read_data,
input_setup,
imsave,
psnr,
merge
)
import time
import os
import cv2
import matplotlib.pyplot as plt
from skimage import data, exposure, img_as_float
import numpy as np
import tensorflow as tf
#定义SRCNN类
class SRCNN(object):
def __init__(self,
sess,
image_size=33,
label_size=21,
batch_size=64,
c_dim=1,
checkpoint_dir=None,
sample_dir=None):
self.sess = sess
self.is_grayscale = (c_dim == 1)
self.image_size = image_size
self.label_size = label_size
self.batch_size = batch_size
self.c_dim = c_dim
self.checkpoint_dir = checkpoint_dir
self.sample_dir = sample_dir
self.build_model()
#搭建网络
def build_model(self):
self.images = tf.placeholder(tf.float32, [None, self.image_size, self.image_size, self.c_dim], name='images')
self.labels = tf.placeholder(tf.float32, [None, self.label_size, self.label_size, self.c_dim], name='labels')
#第一层CNN：对输入图片的特征提取。（9 x 9 x 64卷积核）
#第二层CNN：对第一层提取的特征的非线性映射（1 x 1 x 32卷积核）
#第三层CNN：对映射后的特征进行重建，生成高分辨率图像（5 x 5 x 1卷积核）
#权重
self.weights = {
#论文中为提高训练速度的设置 n1=32 n2=16
'w1': tf.Variable(tf.random_normal([9, 9, 1, 64], stddev=1e-3), name='w1'),
'w2': tf.Variable(tf.random_normal([1, 1, 64, 32], stddev=1e-3), name='w2'),
'w3': tf.Variable(tf.random_normal([5, 5, 32, 1], stddev=1e-3), name='w3')
}
#偏置
self.biases = {
'b1': tf.Variable(tf.zeros([64]), name='b1'),
'b2': tf.Variable(tf.zeros([32]), name='b2'),
'b3': tf.Variable(tf.zeros([1]), name='b3')
}
self.pred = self.model()
# 以MSE作为损耗函数
self.loss = tf.reduce_mean(tf.square(self.labels - self.pred))
self.saver = tf.train.Saver()
#主函数调用(训练或测试)
def train(self, config):
if config.is_train: #判断是否为训练(main传入)
input_setup(self.sess, config)
else:
nx, ny = input_setup(self.sess, config)
#训练为checkpoint下train.h5
#测试为checkpoint下test.h5
if config.is_train:
data_dir = os.path.join('./{}'.format(config.checkpoint_dir), "train.h5")
else:
data_dir = os.path.join('./{}'.format(config.checkpoint_dir), "test.h5")
train_data, train_label = read_data(data_dir)#读取.h5文件(由测试和训练决定)
global_step=tf.Variable(0)#定义global_step 它会自动+1
#通过exponential_decay函数生成学习率
learning_rate_exp=tf.train.exponential_decay(config.learning_rate , global_step , 1480 , 0.98 , staircase=True) #每1个Epoch 学习率*0.98
#标准反向传播的随机梯度下降
#self.train_op = tf.train.GradientDescentOptimizer(config.learning_rate).minimize(self.loss)#学习率learning rate 使self.loss有最小值
self.train_op = tf.train.GradientDescentOptimizer(learning_rate_exp).minimize(self.loss , global_step=global_step)
#Adam 替换上面的连续4行
#self.train_op = tf.train.AdamOptimizer(config.learning_rate).minimize(self.loss, global_step=global_step)
#出现warning : initialize_all_variables (from tensorflow.python.ops.variables) is deprecated and will be removed after 2017-03-02.
#tf.initialize_all_variables().run()
tf.global_variables_initializer().run() #替换掉上句
counter = 0
start_time = time.time()
if self.load(self.checkpoint_dir):
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
#训练
if config.is_train:
print("Training...")
for ep in range(config.epoch): #迭代次数的循环
# 以batch为单元
batch_idxs = len(train_data) // config.batch_size
for idx in range(0, batch_idxs):
batch_images = train_data[idx*config.batch_size : (idx+1)*config.batch_size]
batch_labels = train_label[idx*config.batch_size : (idx+1)*config.batch_size]
counter += 1
_, err = self.sess.run([self.train_op, self.loss], feed_dict={self.images: batch_images, self.labels: batch_labels})
if counter % 10 == 0: #10的倍数step显示
print("Epoch: [%2d], step: [%2d], time: [%4.4f], loss: [%.8f]" \
% ((ep+1), counter, time.time()-start_time, err))
if counter % 500 == 0: #500的倍数step存储
self.save(config.checkpoint_dir, counter)
#测试
else:
print("Testing...")
result = self.pred.eval({self.images: train_data, self.labels: train_label}) # 从test.h中来
result = merge(result, [nx, ny])
result = result.squeeze()#除去size为1的维度
#result= exposure.adjust_gamma(result, 1.07)#调暗一些
image_path = os.path.join(os.getcwd(), config.sample_dir)
image_path = os.path.join(image_path, "MySRCNN.bmp")
imsave(result, image_path)
def model(self):
conv1 = tf.nn.relu(tf.nn.conv2d(self.images, self.weights['w1'], strides=[1,1,1,1], padding='VALID') + self.biases['b1'])
conv2 = tf.nn.relu(tf.nn.conv2d(conv1, self.weights['w2'], strides=[1,1,1,1], padding='VALID') + self.biases['b2'])
conv3 = tf.nn.conv2d(conv2, self.weights['w3'], strides=[1,1,1,1], padding='VALID') + self.biases['b3']
return conv3
def save(self, checkpoint_dir, step):
model_name = "SRCNN.model"
model_dir = "%s_%s" % ("srcnn", self.label_size)
checkpoint_dir = os.path.join(checkpoint_dir, model_dir) #再一次确定路径为 checkpoint->srcnn_21下
if not os.path.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
self.saver.save(self.sess,
os.path.join(checkpoint_dir, model_name), #文件名为SRCNN.model-迭代次数
global_step=step)
def load(self, checkpoint_dir):
print(" [*] Reading checkpoints...")
model_dir = "%s_%s" % ("srcnn", self.label_size)
checkpoint_dir = os.path.join(checkpoint_dir, model_dir) #路径为checkpoint->srcnn_labelsize(21)
#加载路径下的模型(.meta文件保存当前图的结构; .index文件保存当前参数名; .data文件保存当前参数值)
ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
if ckpt and ckpt.model_checkpoint_path:
ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name)) #saver.restore()函数给出model.-n路径后会自动寻找参数名-值文件进行加载
return True
else:
return False