forked from damon-demon/Transform-Based-Tensor-Net
-
Notifications
You must be signed in to change notification settings - Fork 0
/
BuildModel.py
172 lines (110 loc) · 5.41 KB
/
BuildModel.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
# Video Synthesis via Transform-Based Tensor Neural Network
# Yimeng Zhang
# 8/4/2020
import tensorflow as tf
import LoadData as LD
import DefineParam as DP
# Input: Parameters
pixel_w, pixel_h, batchSize, nPhase, nTrainData, nValData, learningRate, nEpoch, nOfModel, ncpkt, trainFile, valFile, testFile, saveDir, modelDir = DP.get_param()
# Model Building
def build_model(phi, restore=False):
Xinput, Xoutput, Phi, PhiT, Yinput, Epoch_num = LD.pre_calculate(phi)
prediction, predictionSymmetric, transField, lambdaStep, softThr = inference_ista(Xinput, Xoutput, Phi, PhiT, Yinput, Epoch_num, reuse=False)
costMean, costSymmetric = compute_cost(prediction, predictionSymmetric, Xoutput)
if restore is False:
costSparsity = compute_sparsity(transField)
costAll = costMean + 0.01*costSymmetric + 0.001*costSparsity
optmAll = tf.train.AdamOptimizer(learning_rate=learningRate).minimize(costAll)
init = tf.global_variables_initializer()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
saver = tf.train.Saver(tf.global_variables(), max_to_keep=100)
sess = tf.Session(config=config)
if restore is False:
sess.run(init)
# saver.restore(sess, '%s/60.cpkt' % modelDir)
else:
saver.restore(sess, '%s/%d.cpkt' % (modelDir, ncpkt))
var = tf.trainable_variables()
for i in var:
print(i)
if restore is False:
return sess, saver, Xinput, Xoutput, Epoch_num, costMean, costSymmetric, costSparsity, optmAll, Yinput, prediction, lambdaStep, softThr, transField
else:
return sess, saver, Xinput, Xoutput, Yinput, Epoch_num, prediction, transField
# Add weight and bias
def add_conv2d_weight(wShape, nOrder):
return tf.get_variable(shape=wShape, initializer=tf.contrib.layers.xavier_initializer_conv2d(), name='weight_%d' % (nOrder))
def ista_block(layerxk, layeryk, Phi, PhiT, Yinput, Epoch_num):
# Parameters
lambdaStep = tf.Variable(0.1, dtype=tf.float32)
softThr = tf.Variable(0.1, dtype=tf.float32)
t = tf.Variable(1, dtype=tf.float32)
convSize1 = 64
convSize2 = 64
convSize3 = 64
filterSize1 = 3
filterSize2 = 3
filterSize3 = 3
# rk <-- yk
rk = tf.subtract(layeryk[-1], tf.scalar_mul(lambdaStep, tf.multiply(PhiT, tf.subtract(tf.multiply(Phi, layeryk[-1]), Yinput))))
# Drk <-- rk
weight0 = add_conv2d_weight([filterSize1, filterSize1, nOfModel, convSize1], 0)
Drk = tf.nn.conv2d(rk, weight0, strides=[1, 1, 1, 1], padding='SAME')
# Transform Module
weight1 = add_conv2d_weight([filterSize2, filterSize2, convSize1, convSize2], 1)
weight2 = add_conv2d_weight([filterSize3, filterSize3, convSize2, convSize3], 2)
Frk = tf.nn.conv2d(Drk, weight1, strides=[1, 1, 1, 1], padding='SAME')
Frk = tf.nn.relu(Frk)
Frk = tf.nn.conv2d(Frk, weight2, strides=[1, 1, 1, 1], padding='SAME')
field = Frk
# Soft-thresholding Module
softFRk = tf.multiply(tf.sign(Frk), tf.nn.relu(tf.subtract(tf.abs(Frk), softThr)))
# Inverse Transform Module
weight3 = add_conv2d_weight([filterSize3, filterSize3, convSize3, convSize2], 3)
weight4 = add_conv2d_weight([filterSize2, filterSize2, convSize2, convSize1], 4)
FFrk = tf.nn.conv2d(softFRk, weight3, strides=[1, 1, 1, 1], padding='SAME')
FFrk = tf.nn.relu(FFrk)
FFrk = tf.nn.conv2d(FFrk, weight4, strides=[1, 1, 1, 1], padding='SAME')
# get Grk from Rk
weight5 = add_conv2d_weight([filterSize1, filterSize1, convSize1, nOfModel], 5)
Grk = tf.nn.conv2d(FFrk, weight5, strides=[1, 1, 1, 1], padding='SAME')
# xk = rk + Grk
xk = tf.add(rk, Grk)
yk = (1 + t)*xk - t*layerxk[-1]
# Symmetric Restriction for Computing Loss Function
sFrk = tf.nn.conv2d(Frk, weight3, strides=[1, 1, 1, 1], padding='SAME')
sFrk = tf.nn.relu(sFrk)
sFrk = tf.nn.conv2d(sFrk, weight4, strides=[1, 1, 1, 1], padding='SAME')
# rk difference
symmetric = sFrk - Drk
return xk, yk, symmetric, field, lambdaStep, softThr
def inference_ista(Xinput, Xoutput, Phi, PhiT, Yinput, Epoch_num, reuse):
layerxk = []
layeryk = []
layerSymmetric = []
transField = []
layerxk.append(Xinput)
layeryk.append(Xinput)
for i in range(nPhase):
with tf.variable_scope('conv_%d' % (i), reuse=reuse):
xk, yk, convSymmetric, field, lambdaStep, softThr = ista_block(layerxk, layeryk, Phi, PhiT, Yinput, Epoch_num)
layerxk.append(xk)
layeryk.append(yk)
layerSymmetric.append(convSymmetric)
transField.append(field)
return layerxk, layerSymmetric, transField, lambdaStep, softThr
# Cost Computation
def compute_cost(prediction, predictionSymmetric, Xoutput):
costMean = tf.reduce_mean(tf.square(prediction[-1][:, :, :, 1] - Xoutput[:, :, :, 1]))
costSymmetric = 0
for k in range(nPhase):
costSymmetric += tf.reduce_mean(tf.square(predictionSymmetric[k]))
return costMean, costSymmetric
# Sparsity Computation
def compute_sparsity(tensor):
costSparsity = 0
for k in range(nPhase):
costSparsity += tf.reduce_mean(tf.abs(tensor[k])) # l1 norm
return costSparsity