WGAN代码解读及实验总结

最新推荐文章于 2024-07-15 00:54:56 发布

CloudCver

最新推荐文章于 2024-07-15 00:54:56 发布

阅读量6.9k

点赞数 3

分类专栏：计算机视觉 Tensorflow

本文链接：https://blog.csdn.net/CLOUD_J/article/details/94392474

版权

计算机视觉同时被 2 个专栏收录

31 篇文章 1 订阅

订阅专栏

Tensorflow

6 篇文章 0 订阅

订阅专栏

GAN作为图像的另一个新领域，本成为21世纪最好的idea。嘿嘿，最近小试牛刀，下载了个WGAN的代码，这里简单分析下，给大家一个参考。
【提示】
本文预计阅读时间5分钟，带灰色底纹的和加粗的为重要部分哦！

（一）WGAN初识

在这里插入图片描述

（二）代码分析

2.1 main struct

打开代码后，它的主要结构如下图所示。
在这里插入图片描述
我们先看一下wgan_conv主函数，打开之后首先直接到最底main的位置，如下

if __name__ == '__main__':

	os.environ['CUDA_VISIBLE_DEVICES'] = '0'
	# the dir of pic generated
	sample_folder = 'Samples/mnist_wgan_conv'
	if not os.path.exists(sample_folder):
		os.makedirs(sample_folder)

	# net param
	generator = G_conv_mnist()
	discriminator = D_conv_mnist()
	# data param
	data = mnist()

	# run
	wgan = WGAN(generator, discriminator, data)
	wgan.train(sample_folder)

这里做几点阐述

1、首先创建了一个目录用来存储你的生成图像，程序会每隔一段时间输出一个图像。
2、搞了三个类，一个generater生成器网络，一个是discriminator判别器类，然后是数据类。
3、又声明一个对象WGAN网络，然后调用它的train函数

OK至此，主函数结构阐述清楚。那此时你会想generater咋定义？discriminator咋定义？
好一个一个看。

2.2 generator

generator是生成器网络，其实就是搭了一个上采样的网络，先将噪声输入一维向量，通过全连接到更多的数据，然后把它展开成二维的图像，这里我们先用的灰度，你也可以改成彩色。然后再上采样，随意搞得，反正最后你要上采样到和你的正样本图像维度一致。如下所示：

class G_conv_mnist(object):
	def __init__(self):
		self.name = 'G_conv_mnist'

	def __call__(self, z):
		with tf.variable_scope(self.name) as scope:
			#step 1 全连接层，把z白噪声变为8*15*128图
			g = tcl.fully_connected(z, 8*15*128, activation_fn = tf.nn.relu, normalizer_fn=tcl.batch_norm,
									weights_initializer=tf.random_normal_initializer(0, 0.02))
			g = tf.reshape(g, (-1, 8, 15, 128)) 
			#step 2 反卷积/上采样 到16*30*64图    4代表卷积核大小
			g = tcl.conv2d_transpose(g, 64, 4,stride=2, 
									activation_fn=tf.nn.relu, normalizer_fn=tcl.batch_norm, padding='SAME', weights_initializer=tf.random_normal_initializer(0, 0.02))
			#step 3 反卷积/上采样 到32*60*1的图，此时和真实手写体的数据是一样的图
			g = tcl.conv2d_transpose(g, 1, 4, stride=2, 
										activation_fn=tf.nn.sigmoid, padding='SAME', weights_initializer=tf.random_normal_initializer(0, 0.02))
			print(g.shape)
			return g
	@property
	def vars(self):
		return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=self.name)

注意：
这里你会看到一个call函数，它是咋用呢？
一个类下面有个call函数，你就可以生成一个对象后，直接把它当成方法用。例如
class G():
call(x):
print(x)
这样的话你就A = G()，然后再A(1)就打印了1。
其实就是说这个类弄好了，之后可以直接当函数用。

好，然后我们看一下discriminator

2.3 discriminator

和generator干了差不多的事情，他要把X和GX输进去，然后搭建一个卷积网络判别真假。

class D_conv_mnist(object):
	def __init__(self):
		self.name = 'D_conv_mnist'
	def __call__(self, x, reuse=False):
		with tf.variable_scope(self.name) as scope:
			if reuse:
				scope.reuse_variables()
			size = 64
			#step 1 卷积4*4卷积核 bzx30x60x1 -> bzx15x30x64
			shared = tcl.conv2d(x, num_outputs=size, kernel_size=4, 
						stride=2, activation_fn=lrelu)
			#step 2 卷积4*4卷积核 bzx15x30x64 -> bzx7x15x128
			shared = tcl.conv2d(shared, num_outputs=size * 2, kernel_size=4,
						stride=2, activation_fn=lrelu, normalizer_fn=tcl.batch_norm)
			#step 3 展开向量 bzx7x15x128 -> bzx6372
			shared = tcl.flatten(shared)
			#step 4 全连接层
			d = tcl.fully_connected(shared, 1, activation_fn=None, weights_initializer=tf.random_normal_initializer(0, 0.02))
			q = tcl.fully_connected(shared, 128, activation_fn=lrelu, normalizer_fn=tcl.batch_norm)
			q = tcl.fully_connected(q, 10, activation_fn=None) # 10 classes
			return d, q
	@property
	def vars(self):
		return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=self.name)

2.4 数据的导入改写

我下载的代码是直接导入的minist数据集，我们可能要导入图片数据集哈，这里我做了一些改变。
这里加了个next_batch函数，先生成随机序列，然后读取batch图像，存到数据集中。

class mnist():
	def __init__(self, flag='conv', is_tanh = False):
		self.datapath = prefix + 'bus_data/'
		self.X_dim = 784 # for mlp
		self.z_dim = 100
		self.y_dim = 10
		self.sizex = 32 # for conv
		self.sizey = 60 # for conv
		self.channel = 1 # for conv
		#self.data = input_data.read_data_sets(datapath, one_hot=True)
		self.flag = flag
		self.is_tanh = is_tanh
		self.Train_nums = 17

	def __call__(self,batch_size):
		batch_imgs = self.next_batch(self.datapath,batch_size)
		#batch_imgs,y = self.next_batch(prefix,batch_size)
		if self.flag == 'conv':
			batch_imgs = np.reshape(batch_imgs, (batch_size, self.sizex, self.sizey, self.channel)) 
		if self.is_tanh:
			batch_imgs = batch_imgs*2 - 1
		#return batch_imgs, y
		return batch_imgs

	def next_batch(self,data_path, batch_size):
	#def next_batch(self,data_path, lable_path, batch_size):
		train_temp = np.random.randint(low=0, high=self.Train_nums, size=batch_size) # 生成元素的值在[low,high)区间，随机选取
		train_data_batch = np.zeros([batch_size,self.sizex,  self.sizey]) # 其中[img_row,  img_col, 3]是原数据的shape，相应变化
		#train_label_batch = np.zeros([batch_size, self.size, self.size]) #
		count = 0 # 后面就是读入图像，并打包成四维的batch
		#print(data_path)
		img_list = os.listdir(data_path)
		#print(img_list)
	
		for i in train_temp:
			img_path = os.path.join(data_path, img_list[i])  # 图片文件
			img = cv.imread(img_path)
			gray = cv.cvtColor(img,cv.COLOR_RGB2GRAY)
			train_data_batch[count, :, :]  = cv.resize(gray,(self.sizey,  self.sizex))
			count+=1
		return train_data_batch#, train_label_batch
	def data2fig(self, samples):
		if self.is_tanh:
			samples = (samples + 1)/2
		fig = plt.figure(figsize=(4, 4))
		gs = gridspec.GridSpec(4, 4)
		gs.update(wspace=0.05, hspace=0.05)

		for i, sample in enumerate(samples):
			ax = plt.subplot(gs[i])
			plt.axis('off')
			ax.set_xticklabels([])
			ax.set_yticklabels([])
			ax.set_aspect('equal')
			plt.imshow(sample.reshape(self.sizex,self.sizey), cmap='Greys_r')
		return fig

2.5 WGAN网络

首先是搭网络NET，discriminator分别把真实的正样本X投进去，把噪声产生的G_sample投进去，得到正负结果。

# nets
		self.G_sample = self.generator(self.z)

		self.D_real, _ = self.discriminator(self.X)
		self.D_fake, _ = self.discriminator(self.G_sample, reuse = True)

然后就是计算损失。我们利用上面结果分别计算D和G的损失，然后有两个优化器，分别对于D和G

# loss
		self.D_loss = - tf.reduce_mean(self.D_real) + tf.reduce_mean(self.D_fake)
		self.G_loss = - tf.reduce_mean(self.D_fake)

		self.D_solver = tf.train.RMSPropOptimizer(learning_rate=1e-4).minimize(self.D_loss, var_list=self.discriminator.vars)
		self.G_solver = tf.train.RMSPropOptimizer(learning_rate=1e-4).minimize(self.G_loss, var_list=self.generator.vars)

这里网络就搭建好了，我们要看一下train函数。其主要是先优化D再优化G这个步骤，这里我么此优化G和D的次数相同，你也可以去调整这个n_d。

for epoch in range(training_epoches):
			# update D
			n_d = 20 if epoch < 250 or (epoch+1) % 500 == 0 else 10
			for _ in range(n_d):
				#X_b, _ = self.data(batch_size)
				X_b= self.data(batch_size)
				self.sess.run(self.clip_D)
				self.sess.run(
						self.D_solver,
            			feed_dict={self.X: X_b, self.z: sample_z(batch_size, self.z_dim)}
						)
			# update G
			for _ in range(n_d):
				#X_b, _ = self.data(batch_size)
				X_b= self.data(batch_size)
				self.sess.run(self.clip_D)
				self.sess.run(
					self.G_solver,
					feed_dict={self.z: sample_z(batch_size, self.z_dim)}
				)

对于WGAN的全部代码如下

class WGAN():
	def __init__(self, generator, discriminator, data):
		self.generator = generator
		self.discriminator = discriminator
		self.data = data

		self.z_dim = self.data.z_dim
		self.sizex = self.data.sizex
		self.sizey = self.data.sizey
		self.channel = self.data.channel

		self.X = tf.placeholder(tf.float32, shape=[None, self.sizex, self.sizey, self.channel])
		self.z = tf.placeholder(tf.float32, shape=[None, self.z_dim])
		# nets
		self.G_sample = self.generator(self.z)

		self.D_real, _ = self.discriminator(self.X)
		self.D_fake, _ = self.discriminator(self.G_sample, reuse = True)

		# loss
		self.D_loss = - tf.reduce_mean(self.D_real) + tf.reduce_mean(self.D_fake)
		self.G_loss = - tf.reduce_mean(self.D_fake)

		self.D_solver = tf.train.RMSPropOptimizer(learning_rate=1e-4).minimize(self.D_loss, var_list=self.discriminator.vars)
		self.G_solver = tf.train.RMSPropOptimizer(learning_rate=1e-4).minimize(self.G_loss, var_list=self.generator.vars)
		
		# clip
		self.clip_D = [var.assign(tf.clip_by_value(var, -0.01, 0.01)) for var in self.discriminator.vars]
		
		gpu_options = tf.GPUOptions(allow_growth=True)
		self.sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

	def train(self, sample_folder, training_epoches = 100000, batch_size = 5):
		i = 0
		self.sess.run(tf.global_variables_initializer())
		
		for epoch in range(training_epoches):
			# update D
			n_d = 20 if epoch < 250 or (epoch+1) % 500 == 0 else 10
			for _ in range(n_d):
				#X_b, _ = self.data(batch_size)
				X_b= self.data(batch_size)
				self.sess.run(self.clip_D)
				self.sess.run(
						self.D_solver,
            			feed_dict={self.X: X_b, self.z: sample_z(batch_size, self.z_dim)}
						)
			# update G
			for _ in range(n_d):
				#X_b, _ = self.data(batch_size)
				X_b= self.data(batch_size)
				self.sess.run(self.clip_D)
				self.sess.run(
					self.G_solver,
					feed_dict={self.z: sample_z(batch_size, self.z_dim)}
				)

			# print loss. save images.
			if epoch % 100 == 0 or epoch < 100:
				D_loss_curr = self.sess.run(
						self.D_loss,
            			feed_dict={self.X: X_b, self.z: sample_z(batch_size, self.z_dim)})
				G_loss_curr = self.sess.run(
						self.G_loss,
						feed_dict={self.z: sample_z(batch_size, self.z_dim)})
				print('Iter: {}; D loss: {:.4}; G_loss: {:.4}'.format(epoch, D_loss_curr, G_loss_curr))

				if epoch % 1000 == 0:
					samples = self.sess.run(self.G_sample, feed_dict={self.z: sample_z(16, self.z_dim)})
					print(samples.shape)
					fig = self.data.data2fig(samples)
					plt.savefig('{}/{}.png'.format(sample_folder, str(i).zfill(3)), bbox_inches='tight')
					i += 1
					plt.close(fig)

（三）实验结果

我找了17张车的图片~~原谅我比较懒，如下图所示。基本都是差不多样子的。
在这里插入图片描述
然后代码跑起来~我们把它resize到（30，60），主要是为了让我的机器跑快些，本来就怂。
一开始是一堆噪声图如下图所示。
其实在训练一段时间后如下所示，可以看出具有一定车的样子，中间黑车身，貌似也能看到个车轱辘。哈哈。初见效果~
在这里插入图片描述
损失的结果如下所示：