/
gain.py
146 lines (109 loc) · 4.76 KB
/
gain.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
import os
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
import numpy as np
from tqdm import tqdm
from utils import normalization, renormalization
from utils import xavier_init
from utils import binary_sampler, uniform_sampler, sample_batch_index
def GAIN (data_x, gain_parameters):
# Define mask matrix
data_m = 1-np.isnan(data_x)
# System parameters
batch_size = gain_parameters['batch_size']
hint_rate = gain_parameters['hint_rate']
alpha = gain_parameters['alpha']
iterations = gain_parameters['iterations']
# Other parameters
no, dim = data_x.shape
# Hidden state dimensions
h_dim = int(dim)
norm_data, norm_parameters = normalization(data_x)
norm_data_x = np.nan_to_num(norm_data, 0)
## GAIN architecture
# Input placeholders
# Data vector
X = tf.placeholder(tf.float32, shape = [None, dim])
# Mask vector
M = tf.placeholder(tf.float32, shape = [None, dim])
# Hint vector
H = tf.placeholder(tf.float32, shape = [None, dim])
Z = tf.placeholder(tf.float32, shape = [None, dim])
# Discriminator variables
D_W1 = tf.Variable(xavier_init([dim*2, h_dim])) # Data + Hint as inputs
D_b1 = tf.Variable(tf.zeros(shape = [h_dim]))
D_W2 = tf.Variable(xavier_init([h_dim, h_dim]))
D_b2 = tf.Variable(tf.zeros(shape = [h_dim]))
D_W3 = tf.Variable(xavier_init([h_dim, dim]))
D_b3 = tf.Variable(tf.zeros(shape = [dim])) # Multi-variate outputs
theta_D = [D_W1, D_W2, D_W3, D_b1, D_b2, D_b3]
#Generator variables
# Data + Mask as inputs (Random noise is in missing components)
G_W1 = tf.Variable(xavier_init([dim*2, h_dim]))
G_b1 = tf.Variable(tf.zeros(shape = [h_dim]))
G_W2 = tf.Variable(xavier_init([h_dim, h_dim]))
G_b2 = tf.Variable(tf.zeros(shape = [h_dim]))
G_W3 = tf.Variable(xavier_init([h_dim, dim]))
G_b3 = tf.Variable(tf.zeros(shape = [dim]))
theta_G = [G_W1, G_W2, G_W3, G_b1, G_b2, G_b3]
## GAIN functions
# Generator
def generator(x, m):
# Concatenate Mask and Data
inputs = tf.concat(values = [x, m], axis = 1)
G_h1 = tf.nn.relu(tf.matmul(inputs, G_W1) + G_b1)
G_h2 = tf.nn.relu(tf.matmul(G_h1, G_W2) + G_b2)
# MinMax normalized output
G_prob = tf.nn.sigmoid(tf.matmul(G_h2, G_W3) + G_b3)
return G_prob
# Discriminator
def discriminator(x, h):
# Concatenate Data and Hint
inputs = tf.concat(values = [x, h], axis = 1)
D_h1 = tf.nn.relu(tf.matmul(inputs, D_W1) + D_b1)
D_h2 = tf.nn.relu(tf.matmul(D_h1, D_W2) + D_b2)
D_logit = tf.matmul(D_h2, D_W3) + D_b3
D_prob = tf.nn.sigmoid(D_logit)
return D_prob, D_logit
## GAIN structure
# Generator
G_sample = generator(X, M)
# Combine with observed data
Hat_X = X * M + G_sample * (1-M)
# Discriminator
D_prob, D_logit = discriminator(Hat_X, H)
## GAIN loss
D_loss_temp = -tf.reduce_mean(M * tf.log(D_prob + 1e-8) + (1-M) * tf.log(1. - D_prob + 1e-8))
G_loss_temp = -tf.reduce_mean((1-M) * tf.log(D_prob + 1e-8))
MSE_loss = tf.reduce_mean((M * X - M * G_sample) * (M * X - M * G_sample)) / tf.reduce_mean(M)
D_loss = D_loss_temp
G_loss = G_loss_temp + alpha * MSE_loss
## GAIN solver
D_solver = tf.train.AdamOptimizer().minimize(D_loss, var_list=theta_D)
G_solver = tf.train.AdamOptimizer().minimize(G_loss, var_list=theta_G)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
## Iterations
for it in tqdm(range(iterations)):
# Sample batch
batch_idx = sample_batch_index(no, batch_size)
X_mb = norm_data_x[batch_idx, :]
M_mb = data_m[batch_idx, :]
# Sample random vectors
Z_mb = uniform_sampler(0, 0.01, batch_size, dim)
# Sample hint vectors
H_mb_temp = binary_sampler(hint_rate, batch_size, dim)
#H_mb = M_mb * H_mb_temp+0.5*(1-H_mb_temp)
H_mb = M_mb * H_mb_temp
# Combine random vectors with observed vectors
X_mb = M_mb * X_mb + (1-M_mb) * Z_mb
_, D_loss_curr, D_logit_curr, D_prob_curr = sess.run([D_solver, D_loss_temp, D_logit, D_prob], feed_dict = {M: M_mb, X: X_mb, H:H_mb})
_, G_loss_curr, MSE_loss_curr = sess.run([G_solver, G_loss_temp, MSE_loss], feed_dict = {X: X_mb, M: M_mb, H:H_mb})
Z_mb = uniform_sampler(0, 0.01, no, dim)
M_mb = data_m
X_mb = norm_data_x
X_mb = M_mb * X_mb + (1-M_mb) * Z_mb
imputed_data = sess.run([G_sample], feed_dict = {X: X_mb, M: M_mb})[0]
imputed_data = data_m * norm_data_x + (1-data_m) * imputed_data
imputed_data= renormalization(imputed_data, norm_parameters)
return imputed_data