from tensorflow import keras import tensorflow as tf from tensorflow.keras import layers import numpy as np import os import matplotlib.pyplot as plt # 设置相关底层配置 physical_devices = tf.config.experimental.list_physical_devices('GPU') assert len(physical_devices) > 0, "Not enough GPU hardware devices available" tf.config.experimental.set_memory_growth(physical_devices[0], True) # os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # os.environ["CUDA_VISIBLE_DEVICES"] = "1" # 使用第2块gpu # 拼接图片 def my_save_img(data,save_path): # 新图拼接行列 r_c = 10 len_data = data.shape[0] each_pix = 64 save_img_path = save_path new_img = np.zeros((r_c*each_pix,r_c*each_pix,3)) for index,each_img in enumerate(data[:r_c*r_c]): # print('each_img.shape:',each_img.shape,np.max(each_img),np.min(each_img)) each_img = (each_img+1)/2 # print('each_img.shape:', each_img.shape, np.max(each_img), np.min(each_img)) row_start = int(index/r_c) * each_pix col_start = (index%r_c)*each_pix # print(index,row_start,col_start) new_img[row_start:row_start+each_pix,col_start:col_start+each_pix,:] = each_img # print('new_img:',new_img) plt.imsave(save_img_path,new_img) class Generator(keras.Model): def __init__(self): super(Generator,self).__init__() # z: [b, 100] => [b, 3*3*512] => [b, 3, 3, 512] => [b, 64, 64, 3] self.fc = layers.Dense(3 * 3 * 512) self.Tconv1 = layers.Conv2DTranspose(256, 3, 3, 'valid') self.bn1 = layers.BatchNormalization() self.Tconv2 = layers.Conv2DTranspose(128, 5, 2, 'valid') self.bn2 = layers.BatchNormalization() self.Tconv3 = layers.Conv2DTranspose(3, 4, 3, 'valid') def call(self, inputs, training=None, mask=None): # [z, 100] => [z, 3*3*512] x = self.fc(inputs) x = tf.reshape(x, [-1, 3, 3, 512]) x = tf.nn.leaky_relu(x) # x = tf.nn.leaky_relu(self.bn1(self.Tconv1(x), training=training)) x = tf.nn.leaky_relu(self.bn2(self.Tconv2(x), training=training)) x = self.Tconv3(x) x = tf.tanh(x) return x class Discriminator(keras.Model): def __init__(self): super(Discriminator,self).__init__() # [b, 64, 64, 3] => [b, 1] self.conv1 = layers.Conv2D(64,5,3,'valid') self.conv2 = layers.Conv2D(128, 5, 3, 'valid') self.bn2 = layers.BatchNormalization() self.conv3 = layers.Conv2D(256, 5, 3, 'valid') self.bn3 = layers.BatchNormalization() # [b,h,w,3] => [b,-1] self.flatten = layers.Flatten() self.fc = layers.Dense(1) def call(self, inputs, training=None, mask=None): x = tf.nn.leaky_relu(self.conv1(inputs)) x = tf.nn.leaky_relu(self.bn2(self.conv2(x),training = training)) x = tf.nn.leaky_relu(self.bn3(self.conv3(x), training=training)) # 打平 x = self.flatten(x) # [b,-1] => [b,1] logits = self.fc(x) return logits def main(): # 超参数 z_dim = 100 epochs = 3000000 batch_size = 1024 learning_rate = 0.002 is_training = True img_data = np.load('img.npy') train_db = tf.data.Dataset.from_tensor_slices(img_data).shuffle(10000).batch(batch_size) sample = next(iter(train_db)) print(sample.shape, tf.reduce_max(sample).numpy(), tf.reduce_min(sample).numpy()) train_db = train_db.repeat() db_iter = iter(train_db) # 判别器 d = Discriminator() # d.build(input_shape=(None, 64, 64, 3)) # 生成器 g = Generator() # g.build(input_shape=(None, z_dim)) # 分别定义优化器 g_optimizer = tf.optimizers.Adam(learning_rate=learning_rate, beta_1=0.5) d_optimizer = tf.optimizers.Adam(learning_rate=learning_rate, beta_1=0.5) for epoch in range(epochs): batch_z = tf.random.uniform([batch_size, z_dim], minval=-1., maxval=1.) batch_x = next(db_iter) # train D with tf.GradientTape() as tape: # 1. treat real image as real # 2. treat generated image as fake fake_image = g(batch_z, is_training) d_fake_logits = d(fake_image, is_training) d_real_logits = d(batch_x, is_training) d_loss_real = tf.nn.sigmoid_cross_entropy_with_logits(logits=d_real_logits,labels=tf.ones_like(d_real_logits)) # d_loss_real = tf.reduce_mean(d_loss_real) d_loss_fake = tf.nn.sigmoid_cross_entropy_with_logits(logits=d_fake_logits,labels=tf.ones_like(d_fake_logits)) # d_loss_fake = tf.reduce_mean(d_loss_fake) d_loss = d_loss_fake + d_loss_real grads = tape.gradient(d_loss, d.trainable_variables) d_optimizer.apply_gradients(zip(grads, d.trainable_variables)) with tf.GradientTape() as tape: fake_image = g(batch_z, is_training) d_fake_logits = d(fake_image, is_training) g_loss = tf.nn.sigmoid_cross_entropy_with_logits(logits=d_fake_logits,labels=tf.ones_like(d_fake_logits)) # g_loss = tf.reduce_mean(g_loss) grads = tape.gradient(g_loss, g.trainable_variables) g_optimizer.apply_gradients(zip(grads, g.trainable_variables)) if epoch % 10 == 0: # print(epoch, 'd-loss:',float(d_loss), 'g-loss:', float(g_loss)) print(epoch, 'd-loss:', d_loss.numpy(), 'g-loss:', g_loss.numpy()) if epoch % 50 == 0: z = tf.random.uniform([225,z_dim]) fake_image = g(z,training = False) img_path = os.path.join('g_pic2', 'gan-%d.png'%epoch) my_save_img(fake_image,img_path) if __name__ == '__main__': main()