[tensorflow] tf2.0 简单例子

tf2.0笔记

感觉，都统一了，pytorch tensorflow mxnet，大家都差不多了

gan例子笔记

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")