感觉，都统一了，pytorch tensorflow mxnet，大家都差不多了
import tensorflow as tf from tensorflow.keras import Model,layers import numpy as np from tensorflow.keras.datasets import mnist num_features = 784 lr_generator = 0.0002 lr_descriminator = 0.0002 training_steps = 20000 batch_size = 128 display_step = 500 noise_dim = 500 def getDataset(): (x_train,y_train),(x_test,y_test) = mnist.load_data() x_train,x_test = np.array(x_train,np.float32),np.array(x_test,np.float32) x_train,x_test = x_train/255.0,x_test/255.0 return x_train,y_train,x_test,y_test x_train,y_train,x_test,y_test = getDataset() # n轴拆分 train_data = tf.data.Dataset.from_tensor_slices((x_train,y_train)) # 这里学习一下 # tf.data.Dataset.repeat(count) 为空或-1无限延长 # shuffle这里填的buffer_size是一个epoch的样本数 # batch化 # 预读取一个数据 train_data = train_data.repeat().shuffle(10000).batch(batch_size).prefetch(1) # Generator 过程 ''' b*500 -(fc之后)-> n*(7*7*128) -(reshape)-> n*7*7*128 -(upsample)-> n*14*14*64 -(upsample)->n*28*28*1 ''' class Generator(Model): def __init__(self): super(Generator,self).__init__() self.fc1 = layers.Dense(7*7*128) self.bn1 = layers.BatchNormalization() # upsample卷积，反卷积 # https://github.com/vdumoulin/conv_arithmetic/raw/master/gif/padding_strides_transposed.gif # 洞洞卷积，相当于same的stride=1，w=14，所以输出14*14 self.conv2tr1 = layers.Conv2DTranspose(64,5,strides=2,padding="SAME")#filters,kernel size # 在batch维度和channel维度标准化 self.bn2 = layers.BatchNormalization() self.conv2tr2 = layers.Conv2DTranspose(1,5,strides=2,padding="SAME") def __call__(self,x,is_training = False): x = self.fc1(x) x = self.bn1(x,training = is_training) # leaky_relu x<0时为x/a而不是0，防止梯度消失 x = tf.nn.leaky_relu(x) x = tf.reshape(x,shape = [-1,7,7,128]) x = self.conv2tr1(x) x = self.bn2(x,training = is_training) x = tf.nn.leaky_relu(x) x = self.conv2tr2(x) x = tf.nn.tanh(x) return x # Discriminator 过程 ''' n*768 -> n*28*28*1 -> n*14*14*64 -> n*7*7*128 -> n*(7*7*128) -> n*1024 -> n*2 ''' class Discriminator(Model): def __init__(self): super(Discriminator,self).__init__() self.conv1 = layers.Conv2D(64,5,strides = 2,padding = "SAME") self.bn1 = layers.BatchNormalization() self.conv2 = layers.Conv2D(128,5,strides = 2,padding = "SAME") self.bn2 = layers.BatchNormalization() self.flatten = layers.Flatten() self.fc1 = layers.Dense(1024) self.bn3 = layers.BatchNormalization() self.fc2 = layers.Dense(2) def __call__(self,is_training = False): x = tf.reshape(x,[-1,28,28,1]) x = self.conv1(x) x = self.bn1(x,training = is_training) x = tf.nn.leaky_relu(x) x = self.conv2(x) x = self.bn2(x,training = is_training) x = tf.nn.leaky_relu(x) x = self.flatten(x) x = self.fc1(x) x = self.bn3(x,training = is_training) x = tf.nn.leaky_relu() x = self.fc2(x) return x generator = Generator() discriminator = Discriminator() def generator_loss(reconstructed_image): gen_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits = reconstructed_image,labels = tf.ones([batch_size],dtype = tf.int32))) return gen_loss def discriminator_loss(disc_fake,disc_real): # loss、 disc_loss_real = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=disc_real,labels = tf.ones([batch_size],dtype=tf.int32))) disc_loss_fake = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits = disc_fake,labels = tf.zeros([batch_size],dtype = tf.int32))) return disc_loss_real + disc_loss_fake optimizer_gen = tf.optimizers.Adam(learning_rate = lr_generator) optimizer_disc = tf.optimizers.Adam(learning_rate = lr_descriminator) def run_optimization(real_images): real_images = real_images * 2. -1 #(-1,1)范围内 noise = np.random.normal(-1.,1.,size=[batch_size,noise_dim]) # 通过随机生成噪声数据，用正太分布的噪声去生成图片，生成器的作用就是生成fake images with tf.GradientTape() as g: fake_images = generator(noise,is_training = True) disc_fake = discriminator(fake_images,is_training = True) disc_real = discriminator(real_images,is_training = True) disc_loss = discriminator_loss(disc_fake,disc_real) gradients_disc = g.gradient(disc_loss,discriminator.trainable_vatiables) optimizer_disc.apply_gradients(zip(gradients_disc,discriminator.trainable_variables)) # 由于上面判别器的梯度已经进行更新了,这里又用到判别器来判别fake_images，上面会影响这里判别器的判断，所以不能直接用前面生成好的噪声数据 # 我认为判别器梯度更新在前应该有利于收敛吧，不然最开始先更新生成器梯度的话，最开始训练的时候效果应该不太好 noise = np.random.normal(-1.,1.,size = [batch_size,noise_dim]).astype(np.float32) with tf.GradientTape() as g: fake_images = generator(noise,is_training = True) disc_fake = discriminator(fake_images) gen_loss = generator_loss(disc_fake) gradients_gen = g.gradient(gen_loss,generator.trainable_variables) optimizer_gen.apply_gradients(zip(gradients_gen, generator.trainable_variables)) return gen_loss,disc_loss for step, (batch_x, _) in enumerate(train_data.take(training_steps + 1)): if step == 0: noise = np.random.normal(-1., 1., size=[batch_size, noise_dim]).astype(np.float32) gen_loss = generator_loss(discriminator(generator(noise))) disc_loss = discriminator_loss(discriminator(batch_x), discriminator(generator(noise))) print("initial: gen_loss: %f, disc_loss: %f" % (gen_loss, disc_loss)) continue gen_loss, disc_loss = run_optimization(batch_x) if step % display_step == 0: print("step: %i, gen_loss: %f, disc_loss: %f" % (step, gen_loss, disc_loss)) # 保存权重 generator.save_weights(file_path = "./gen.ckpt") discriminator.save_weights(file_path = "./disc.ckpt")