Skip to content

O
OneFlow-Benchmark

Oneflow-Inc / OneFlow-Benchmark
上一次同步 2 年多

通知 1
Star 92
Fork 0
O
OneFlow-Benchmark

前往新版Gitcode，体验更适合开发者的 AI 搜索 >>

未验证 提交 8b9075ab 编写于 作者: J JamiePlur 提交者： GitHub
浏览文件
操作

Merge pull request #71 from Oneflow-Inc/release_gan 

Release gan
上级 23e573b8 51d0ca93
隐藏空白更改
内联 并排
Showing with 322 addition and 380 deletion
Generative/README.md
浏览文件 @ 8b9075ab
......@@ -28,17 +28,16 @@ python dcgan.py
- `-lr` 学习率，默认为1e-4
- `-e` 设置训练epoch次数，默认为10
- `-c` 与tensorflow进行对齐测试
- `-b` 设置batchsize
- `-g` 设置gpu数量
- `-m` 多机训练
其他需要注意的是：
- 训练将会默认使用minst数据集，如果第一次使用脚本，将会默认将数据集下载到`.data/`目录
- 训练结束后，将会默认保存模型到`.checkpoint/`目录下
- 模型的结构和参数参考了tensorflow的[官方示例](https://www.tensorflow.org/tutorials/generative/dcgan)，可以通过-c参数来跟tensorflow的实现进行对齐测试
- 模型的结构和参数参考了tensorflow的[官方示例](https://www.tensorflow.org/tutorials/generative/dcgan)
- 模型会定期将生成的图片存储到`.gout/`目录，并在训练结束后生成图片演化的动图,生成动图的过程会依赖python包`imageio`
![dcgan demo](https://github.com/Oneflow-Inc/OneFlow-Benchmark/blob/dev_gan/Generative/pic/1.png)
......@@ -57,10 +56,9 @@ python pix2pix.py
- `-lr` 学习率，默认为1e-4
- `-e` 设置训练epoch次数，默认为10
- `-c` 进行对齐测试
- `-b` 设置batchsize
- `-g` 设置gpu数量
- `-m` 多机训练
......@@ -68,7 +66,7 @@ python pix2pix.py
- 训练将会默认使用CMP Facade数据集，如果第一次使用脚本，将会默认将数据集下载到`.data/`目录
- 训练结束后，将会默认保存模型到`.checkpoint/`目录下
- 模型的结构和参数参考了tensorflow的[官方示例](https://www.tensorflow.org/tutorials/generative/pix2pix)，可以通过-c参数来跟tensorflow的实现进行对齐测试
- 模型的结构和参数参考了tensorflow的[官方示例](https://www.tensorflow.org/tutorials/generative/pix2pix)
- 模型会在训练中定期将生成的图片存储到`.gout/`目录
![pix2pix demo](https://github.com/Oneflow-Inc/OneFlow-Benchmark/blob/dev_gan/Generative/pic/2.png)
Generative/dcgan.py
浏览文件 @ 8b9075ab
......@@ -11,80 +11,33 @@ class DCGAN:
def __init__(self, args):
self.lr = args.learning_rate
self.z_dim = 100
self.eval_interval = 400
self.eval_interval = 100
self.eval_size = 16
self.seed = np.random.normal(0, 1, size=(self.eval_size, self.z_dim)).astype(
np.float32
)
self.gpus_per_node = args.gpu_num_per_node
self.multi_nodes = args.multinode
self.batch_size = args.batch_size * self.gpus_per_node
def compare_with_tf(self):
from tf_dcgan import tf_dcgan_test
tf_dcgan_test()
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
func_config.default_distribute_strategy(flow.distribute.consistent_strategy())
func_config.train.primary_lr(self.lr)
func_config.train.model_update_conf(dict(naive_conf={}))
@flow.function(func_config)
def CompareJob(
z=flow.FixedTensorDef((self.batch_size, 100)),
label1=flow.FixedTensorDef((self.batch_size, 1)),
):
img = self.gennerator(z, const_init=True)
logit = self.discriminator(img, const_init=True)
loss = flow.nn.sigmoid_cross_entropy_with_logits(label1, logit)
flow.losses.add_loss(loss)
return loss
check_point = flow.train.CheckPoint()
check_point.init()
z = np.load("z.npy")
label1 = np.ones((self.batch_size, 1)).astype(np.float32)
label0 = np.zeros((self.batch_size, 1)).astype(np.float32)
of_out = CompareJob(z, label1).get()
tf_out = np.load("out.npy")
assert np.allclose(of_out.ndarray(), tf_out, rtol=1e-2, atol=1e-2)
print("compare test passed")
def init_nodes(self):
flow.env.ctrl_port(12138)
nodes = []
for ip in ['192.168.1.12', '192.168.1.14']:
addr_dict = {}
addr_dict["addr"] = ip
nodes.append(addr_dict)
flow.env.machine(nodes)
def train(self, epochs=1, model_dir=None, save=True):
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
func_config.default_distribute_strategy(flow.distribute.consistent_strategy())
func_config.train.primary_lr(self.lr)
func_config.train.model_update_conf(dict(adam_conf={"beta1": 0.5}))
# func_config.train.model_update_conf(dict(naive_conf={}))
func_config.train.model_update_conf(dict(adam_conf={}))
flow.config.gpu_device_num(self.gpus_per_node)
if self.multi_nodes:
self.init_nodes()
@flow.global_function(func_config)
def train_generator(
z=flow.FixedTensorDef((self.batch_size, self.z_dim)),
label1=flow.FixedTensorDef((self.batch_size, 1)),
):
g_out = self.generator(z, trainable=True)
g_logits = self.discriminator(g_out, trainable=False)
g_loss = flow.nn.sigmoid_cross_entropy_with_logits(
label1, g_logits, name="Gloss_sigmoid_cross_entropy_with_logits"
flow.ones_like(g_logits), g_logits, name="Gloss_sigmoid_cross_entropy_with_logits"
)
g_loss = flow.math.reduce_mean(g_loss)
flow.losses.add_loss(g_loss)
return g_loss, g_out
......@@ -92,21 +45,18 @@ class DCGAN:
def train_discriminator(
z=flow.FixedTensorDef((self.batch_size, 100)),
images=flow.FixedTensorDef((self.batch_size, 1, 28, 28)),
label1=flow.FixedTensorDef((self.batch_size, 1)),
label0=flow.FixedTensorDef((self.batch_size, 1)),
):
g_out = self.generator(z, trainable=False)
g_logits = self.discriminator(g_out, trainable=True)
d_loss_fake = flow.nn.sigmoid_cross_entropy_with_logits(
label0, g_logits, name="Dloss_fake_sigmoid_cross_entropy_with_logits"
flow.zeros_like(g_logits), g_logits, name="Dloss_fake_sigmoid_cross_entropy_with_logits"
)
d_logits = self.discriminator(images, trainable=True, reuse=True)
d_loss_real = flow.nn.sigmoid_cross_entropy_with_logits(
label1, d_logits, name="Dloss_real_sigmoid_cross_entropy_with_logits"
flow.ones_like(d_logits), d_logits, name="Dloss_real_sigmoid_cross_entropy_with_logits"
)
d_loss = d_loss_fake + d_loss_real
d_loss = flow.math.reduce_mean(d_loss)
flow.losses.add_loss(d_loss)
return d_loss, d_loss_fake, d_loss_real
......@@ -131,13 +81,11 @@ class DCGAN:
z = np.random.normal(0, 1, size=(self.batch_size, self.z_dim)).astype(
np.float32
)
label1 = np.ones((self.batch_size, 1)).astype(np.float32)
label0 = np.zeros((self.batch_size, 1)).astype(np.float32)
images = x[
batch_idx * self.batch_size : (batch_idx + 1) * self.batch_size
].astype(np.float32)
d_loss, _, _ = train_discriminator(z, images, label1, label0).get()
g_loss, _ = train_generator(z, label1).get()
d_loss, _, _ = train_discriminator(z, images).get()
g_loss, gout = train_generator(z).get()
batch_total = batch_idx + epoch_idx * batch_num * self.batch_size
if (batch_idx + 1) % self.eval_interval == 0:
......@@ -146,9 +94,10 @@ class DCGAN:
epoch_idx + 1, batch_idx + 1, d_loss.mean(), g_loss.mean()
)
)
self._eval_model_and_save_images(
eval_generator, batch_idx + 1, epoch_idx + 1
)
self._save_images(gout, batch_idx + 1, epoch_idx + 1)
# self._eval_model_and_save_images(
# eval_generator, batch_idx + 1, epoch_idx + 1
# )
if save:
from datetime import datetime
if not os.path.exists("checkpoint"):
......@@ -159,6 +108,17 @@ class DCGAN:
)
)
def _save_images(self, results, batch_idx, epoch_idx):
fig = plt.figure(figsize=(4, 4))
for i in range(self.eval_size):
plt.subplot(4, 4, i + 1)
plt.imshow(results[i, 0, :, :] * 127.5 + 127.5, cmap="gray")
plt.axis("off")
if not os.path.exists("gout"):
os.mkdir("gout")
plt.savefig("gout/image_{:02d}_{:04d}.png".format(epoch_idx, batch_idx))
plt.close()
def save_to_gif(self):
anim_file = "dcgan.gif"
with imageio.get_writer(anim_file, mode="I") as writer:
......@@ -331,7 +291,7 @@ class DCGAN:
if transpose:
X = np.transpose(X, (0, 3, 1, 2))
return X / 255.0, y_vec
return (X - 127.5) / 127.5, y_vec
if __name__ == "__main__":
......@@ -340,23 +300,15 @@ if __name__ == "__main__":
parser = argparse.ArgumentParser(description="flags for multi-node and resource")
parser.add_argument("-g", "--gpu_num_per_node", type=int, default=1, required=False)
parser.add_argument("-e", "--epoch_num", type=int, default=10, required=False)
parser.add_argument("-lr", "--learning_rate", type=float, default=1e-4, required=False)
parser.add_argument(
"-c", "--compare", default=False, action="store_true", required=False
)
parser.add_argument(
"-m", "--multinode", default=False, action="store_true", required=False
)
parser.add_argument("-lr", "--learning_rate", type=float, default=1e-5, required=False)
parser.add_argument(
"-load", "--model_load_dir", type=str, default="checkpoint", required=False
)
parser.add_argument(
"-save", "--model_save_dir", type=str, default="checkpoint", required=False
)
parser.add_argument("-b", "--batch_size", type=int, default=8, required=False)
parser.add_argument("-b", "--batch_size", type=int, default=32, required=False)
args = parser.parse_args()
dcgan = DCGAN(args)
if args.compare:
dcgan.compare_with_tf()
dcgan.train(args.epoch_num)
dcgan.save_to_gif()
Generative/layers.py
浏览文件 @ 8b9075ab
import oneflow as flow
import oneflow.python.framework.distribute as distribute_util
def get_const_initializer():
return flow.constant_initializer(0.002)
def deconv2d(
input,
......@@ -28,7 +31,7 @@ def deconv2d(
dtype=input.dtype,
initializer=flow.random_normal_initializer(stddev=0.02)
if not const_init
else flow.constant_initializer(0.002),
else get_const_initializer(),
trainable=trainable,
)
......@@ -77,8 +80,9 @@ def conv2d(
dtype=input.dtype,
initializer=flow.random_normal_initializer(stddev=0.02)
if not const_init
else flow.constant_initializer(0.002),
else get_const_initializer(),
trainable=trainable,
reuse=reuse,
)
output = flow.nn.compat_conv2d(
......@@ -97,15 +101,115 @@ def conv2d(
dtype=input.dtype,
initializer=flow.constant_initializer(0.0),
trainable=trainable,
reuse=reuse,
)
output = flow.nn.bias_add(output, bias, "NCHW")
return output
def batchnorm(input, name, axis=1, reuse=False):
return flow.layers.batch_normalization(input, axis=axis)
# def batchnorm(input, name, axis=1, reuse=False):
# name_ = name if reuse == False else name + "_reuse"
# return flow.layers.batch_normalization(input, axis=axis, name=name_)
def batchnorm(
inputs,
axis=1,
momentum: float = 0.99,
epsilon: float = 0.001,
center: bool = True,
scale: bool = True,
beta_initializer=None,
gamma_initializer=None,
beta_regularizer=None,
gamma_regularizer=None,
moving_mean_initializer=None,
moving_variance_initializer=None,
trainable=True,
training=True,
name: str = "BatchNorm",
reuse=False,
):
name_ = name if reuse == False else name + "_reuse"
if axis < 0:
axis += len(inputs.shape)
assert axis >= 0 and axis < len(inputs.shape)
params_shape = [inputs.shape[axis]]
# Float32 required to avoid precision-loss when using fp16 input/output
params_dtype = flow.float32 if inputs.dtype == flow.float16 else inputs.dtype
if not flow.current_global_function_desc().IsTrainable() or not trainable:
training = False
if center:
beta = flow.get_variable(
name=name + "beta",
shape=params_shape,
dtype=params_dtype,
initializer=beta_initializer or flow.zeros_initializer(),
regularizer=beta_regularizer,
trainable=trainable,
distribute=distribute_util.broadcast(),
reuse=reuse,
)
else:
beta = flow.constant(0, dtype=params_dtype, shape=params_shape, name="beta")
if scale:
gamma = flow.get_variable(
name=name + "gamma",
shape=params_shape,
dtype=params_dtype,
initializer=gamma_initializer or flow.ones_initializer(),
regularizer=gamma_regularizer,
trainable=trainable,
distribute=distribute_util.broadcast(),
reuse=reuse,
)
else:
gamma = flow.constant(
1, dtype=params_dtype, shape=params_shape, name="gamma"
)
moving_mean = flow.get_variable(
name=name + "moving_mean",
shape=params_shape,
dtype=params_dtype,
initializer=moving_mean_initializer or flow.zeros_initializer(),
trainable=False,
distribute=distribute_util.broadcast(),
reuse=reuse,
)
moving_variance = flow.get_variable(
name=name + "moving_variance",
shape=params_shape,
dtype=params_dtype,
initializer=moving_variance_initializer or flow.ones_initializer(),
trainable=False,
distribute=distribute_util.broadcast(),
reuse=reuse,
)
builder = (
flow.user_op_builder(name_)
.Op("normalization")
.Input("x", [inputs])
.Input("moving_mean", [moving_mean])
.Input("moving_variance", [moving_variance])
.Input("gamma", [gamma])
.Input("beta", [beta])
.Output("y")
.Attr("axis", axis)
.Attr("epsilon", epsilon)
.Attr("training", training)
.Attr("momentum", momentum)
)
if trainable and training:
builder = builder.Output("mean").Output("inv_variance")
return builder.Build().InferAndTryRun().RemoteBlobList()[0]
def dense(
input, units, name, use_bias=False, trainable=True, reuse=False, const_init=False
......@@ -124,8 +228,9 @@ def dense(
dtype=inputs.dtype,
initializer=flow.random_normal_initializer(stddev=0.02)
if not const_init
else flow.constant_initializer(0.002),
else get_const_initializer(),
trainable=trainable,
reuse=reuse,
model_name="weight",
)
......@@ -138,8 +243,9 @@ def dense(
dtype=inputs.dtype,
initializer=flow.random_normal_initializer()
if not const_init
else flow.constant_initializer(0.002),
else get_const_initializer(),
trainable=trainable,
reuse=reuse,
model_name="bias",
)
out = flow.nn.bias_add(out, bias, name=name_ + "_bias_add")
......
Generative/pix2pix.py
浏览文件 @ 8b9075ab
......@@ -10,47 +10,73 @@ class Pix2Pix:
self.lr = args.learning_rate
self.out_channels = 3
self.img_size = 256
self.LAMBDA = 50
self.eval_interval = 10
self.LAMBDA = 100
self.eval_interval = 100
self.data_dir = "data/facades/"
self.gpus_per_node = args.gpu_num_per_node
self.batch_size = args.batch_size * self.gpus_per_node
self.multi_nodes = args.multinode
def compare_with_tf(self):
from tf_pix2pix import tf_pix2pix_test
tf_pix2pix_test()
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
func_config.default_distribute_strategy(flow.distribute.consistent_strategy())
func_config.default_distribute_strategy(flow.scope.consistent_view())
func_config.train.primary_lr(1e-4)
func_config.train.model_update_conf(dict(naive_conf={}))
@flow.global_function(func_config)
def CompareJob(
def test_generator(
input=flow.FixedTensorDef((self.batch_size, 3, 256, 256)),
target=flow.FixedTensorDef((self.batch_size, 3, 256, 256)),
):
out = self.generator(input, const_init=True)
out = self.discriminator(out, target, const_init=True)
flow.losses.add_loss(out)
return out
g_out = self.generator(input, trainable=True, const_init=True)
g_logits = self.discriminator(input, g_out, trainable=False, const_init=True)
gan_loss = flow.nn.sigmoid_cross_entropy_with_logits(
flow.ones_like(g_logits), g_logits, name="Gloss_sigmoid_cross_entropy_with_logits"
)
l1_loss = flow.math.reduce_mean(flow.math.abs(g_out - target))
g_loss = gan_loss + self.LAMBDA * l1_loss
flow.losses.add_loss(g_out)
return g_loss
@flow.global_function(func_config)
def test_discriminator(
input=flow.FixedTensorDef((self.batch_size, 3, 256, 256)),
target=flow.FixedTensorDef((self.batch_size, 3, 256, 256)),
):
g_out = self.generator(input, trainable=False, const_init=True)
d_fake_logits = self.discriminator(input, g_out, trainable=True, const_init=True)
d_real_logits = self.discriminator(input, target, trainable=True, reuse=True, const_init=True)
d_fake_loss = flow.nn.sigmoid_cross_entropy_with_logits(
flow.zeros_like(d_fake_logits), d_fake_logits, name="Dloss_fake_sigmoid_cross_entropy_with_logits"
)
d_real_loss = flow.nn.sigmoid_cross_entropy_with_logits(
flow.ones_like(d_real_logits), d_real_logits, name="Dloss_real_sigmoid_cross_entropy_with_logits"
)
d_loss = d_fake_loss + d_real_loss
flow.losses.add_loss(d_loss)
return d_loss
check_point = flow.train.CheckPoint()
check_point.init()
inp = np.load("input.npy").transpose(0, 3, 1, 2).astype(np.float32, order="C")
target = (
np.load("target.npy").transpose(0, 3, 1, 2).astype(np.float32, order="C")
)
of_out = CompareJob(inp, target).get()
tf_out = (
np.load("result.npy").transpose(0, 3, 1, 2).astype(np.float32, order="C")
)
print((tf_out - of_out).mean())
assert np.allclose(of_out.ndarray(), tf_out, rtol=1e-2, atol=1e-2)
tar = np.load("target.npy").transpose(0, 3, 1, 2).astype(np.float32, order="C")
for i in range(3):
d_loss = test_discriminator(inp, tar).get()
g_loss = test_generator(inp, tar).get()
# print("tf produce d_loss:{}, g_loss:{}".format(d_loss.numpy().mean(), g_loss.numpy().mean()))
print("of produce g_loss:{}".format(g_loss.numpy().mean()))
# tf_g_loss = np.load("g_loss.npy")
# tf_d_loss = np.load("d_loss.npy")#.transpose(0, 3, 1, 2).astype(np.float32, order="C")
# assert np.allclose(g_loss.numpy(), tf_g_loss, rtol=1e-1, atol=1e-1), '{}-{}'.format(g_loss.numpy().mean(), tf_g_loss.mean())
# print('G test passed')
# assert np.allclose(d_loss.numpy(), tf_d_loss, rtol=1e-2, atol=1e-2), '{}-{}'.format(d_loss.numpy().mean(), tf_d_loss.mean())
# print('D test passed')
def _downsample(
self,
......@@ -61,7 +87,7 @@ class Pix2Pix:
reuse=False,
apply_batchnorm=True,
trainable=True,
const_init=True,
const_init=False,
):
out = layers.conv2d(
inputs,
......@@ -74,7 +100,7 @@ class Pix2Pix:
name=name + "_conv",
)
if apply_batchnorm:
if apply_batchnorm: #and not const_init:
out = layers.batchnorm(out, name=name + "_bn", reuse=reuse)
out = flow.nn.leaky_relu(out, alpha=0.3)
......@@ -88,7 +114,8 @@ class Pix2Pix:
name,
apply_dropout=False,
trainable=True,
const_init=True,
const_init=False,
reuse = False
):
out = layers.deconv2d(
......@@ -101,19 +128,17 @@ class Pix2Pix:
name=name + "_deconv",
)
# if not const_init:
out = layers.batchnorm(out, name=name + "_bn")
if apply_dropout:
if apply_dropout and not const_init:
out = flow.nn.dropout(out, rate=0.5)
out = flow.keras.activations.relu(out)
out = flow.math.relu(out)
return out
def generator(self, inputs, trainable=True, const_init=False):
if const_init:
apply_dropout = False
else:
apply_dropout = True
apply_dropout = False if const_init else True
# (n, 64, 128, 128)
d1 = self._downsample(
inputs,
......@@ -214,7 +239,7 @@ class Pix2Pix:
const_init=const_init,
trainable=trainable,
)
u0 = flow.keras.activations.tanh(u0)
u0 = flow.math.tanh(u0)
return u0
......@@ -256,8 +281,8 @@ class Pix2Pix:
const_init=const_init,
use_bias=False,
)
bn1 = layers.batchnorm(conv1, name="d_bn1", reuse=reuse)
leaky_relu = flow.nn.leaky_relu(bn1, alpha=0.3)
bn = layers.batchnorm(conv1, name="d_bn", reuse=reuse)
leaky_relu = flow.nn.leaky_relu(bn, alpha=0.3)
# (n, 512, 33, 33)
pad2 = flow.pad(leaky_relu, [[0, 0], [0, 0], [1, 1], [1, 1]])
# (n, 1, 30, 30)
......@@ -275,25 +300,36 @@ class Pix2Pix:
return conv2
def load_facades(self, mode="train"):
from PIL import Image
if not os.path.exists("data/facades/"):
def try_download_facades(self):
if not os.path.exists("data/facades"):
print("not Found Facades - start download")
import tensorflow as tf
if not os.path.exists("data"):
os.mkdir("data")
_PATH = os.path.join(os.getcwd(), "data/facades.tar.gz")
_URL = "https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz"
path_to_zip = tf.keras.utils.get_file(_PATH, origin=_URL, extract=True)
import subprocess
if not os.path.exists("data1"):
os.mkdir("data1")
url = "https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz"
out_path = os.path.join("data1", "facades.tar.gz")
cmd = ["curl", url, "-o", out_path]
print("Downloading ", "facades.tar.gz")
subprocess.call(cmd)
cmd = ["tar", "-xvf", out_path, "-C", "data1"]
print("Decompressing ", "facades.tar.gz")
subprocess.call(cmd)
else:
print("Found Facades - skip download")
def load_facades(self, mode="train"):
from PIL import Image
input_imgs, real_imgs = [], []
for d in os.listdir(os.path.join("data/facades/", mode)):
d = os.path.join("data/facades/", mode, d)
img = np.asarray(Image.open(d))
real_img = img[:, :256, :]
input_img = img[:, 256:, :]
img = Image.open(d)
r1, r2 = np.random.randint(30, size=2)
real_img = img.crop((0, 0, 256, 256)).resize((286, 286))
real_img = np.asarray(real_img.crop((r1, r2, r1 + 256, r2 + 256)))
assert real_img.shape == (256, 256, 3), real_img.shape
input_img = img.crop((256, 0, 512, 256)).resize((286, 286))
input_img = np.asarray(input_img.crop((r1, r2, r1 + 256, r2 + 256)))
assert input_img.shape == (256, 256, 3), input_img.shape
input_imgs.append(input_img)
real_imgs.append(real_img)
......@@ -301,10 +337,17 @@ class Pix2Pix:
real_imgs = np.array(real_imgs).transpose(0, 3, 1, 2)
input_imgs = input_imgs / 127.5 - 1
real_imgs = real_imgs / 127.5 - 1
seed = 547
np.random.seed(seed)
np.random.shuffle(real_imgs)
np.random.seed(seed)
np.random.shuffle(input_imgs)
return input_imgs, real_imgs
def _eval_model(self, model, batch_idx, epoch_idx):
ind = 3
ind = 5
input, target = self.load_facades(mode="test")
gout = model(input[ind * self.batch_size : (ind + 1) * self.batch_size].astype(np.float32, order="C")).get()
plt.figure(figsize=(15, 15))
......@@ -323,50 +366,25 @@ class Pix2Pix:
plt.savefig("gout/image_{:02d}_{:04d}.png".format(epoch_idx + 1, batch_idx + 1))
plt.close()
def test(self):
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
func_config.default_distribute_strategy(flow.distribute.consistent_strategy())
@flow.global_function(func_config)
def eval_generator(input=flow.FixedTensorDef((self.batch_size, 3, 256, 256)),):
g_out = self.generator(input, trainable=False)
return g_out
batch_idx = 1
epoch_idx = 1
check_point = flow.train.CheckPoint()
check_point.init()
x, y = self.load_facades()
batch_num = len(x) // self.batch_size
inp = x[batch_idx * self.batch_size : (batch_idx + 1) * self.batch_size].astype(
np.float32, order="C"
)
target = y[
batch_idx * self.batch_size : (batch_idx + 1) * self.batch_size
].astype(np.float32, order="C")
self._eval_model(eval_generator, batch_idx, epoch_idx)
def train(self, epochs=1, save=True):
self.try_download_facades()
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
func_config.default_distribute_strategy(flow.distribute.consistent_strategy())
func_config.default_distribute_strategy(flow.scope.consistent_view())
func_config.train.primary_lr(self.lr)
func_config.train.model_update_conf(dict(adam_conf={"beta1": 0.5}))
# func_config.train.model_update_conf(dict(adam_conf={"beta1": 0.5}))
func_config.train.model_update_conf(dict(naive_conf={}))
flow.config.gpu_device_num(self.gpus_per_node)
if self.multi_nodes:
self.init_nodes()
@flow.global_function(func_config)
def train_generator(
input=flow.FixedTensorDef((self.batch_size, 3, 256, 256)),
target=flow.FixedTensorDef((self.batch_size, 3, 256, 256)),
label1=flow.FixedTensorDef((self.batch_size, 1, 30, 30)),
):
g_out = self.generator(input, trainable=True)
g_logits = self.discriminator(input, g_out, trainable=False)
gan_loss = flow.nn.sigmoid_cross_entropy_with_logits(
label1, g_logits, name="Gloss_sigmoid_cross_entropy_with_logits"
flow.ones_like(g_logits), g_logits, name="Gloss_sigmoid_cross_entropy_with_logits"
)
l1_loss = flow.math.reduce_mean(flow.math.abs(g_out - target))
g_loss = gan_loss + self.LAMBDA * l1_loss
......@@ -378,18 +396,16 @@ class Pix2Pix:
def train_discriminator(
input=flow.FixedTensorDef((self.batch_size, 3, 256, 256)),
target=flow.FixedTensorDef((self.batch_size, 3, 256, 256)),
label0=flow.FixedTensorDef((self.batch_size, 1, 30, 30)),
label1=flow.FixedTensorDef((self.batch_size, 1, 30, 30)),
):
g_out = self.generator(input, trainable=False)
g_logits = self.discriminator(g_out, target, trainable=True)
g_logits = self.discriminator(input, g_out, trainable=True)
d_fake_loss = flow.nn.sigmoid_cross_entropy_with_logits(
label0, g_logits, name="Dloss_fake_sigmoid_cross_entropy_with_logits"
flow.zeros_like(g_logits), g_logits, name="Dloss_fake_sigmoid_cross_entropy_with_logits"
)
d_logits = self.discriminator(input, target, trainable=True, reuse=True)
d_real_loss = flow.nn.sigmoid_cross_entropy_with_logits(
label1, d_logits, name="Dloss_real_sigmoid_cross_entropy_with_logits"
flow.ones_like(d_logits), d_logits, name="Dloss_real_sigmoid_cross_entropy_with_logits"
)
d_loss = d_fake_loss + d_real_loss
......@@ -399,32 +415,27 @@ class Pix2Pix:
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
func_config.default_distribute_strategy(flow.distribute.consistent_strategy())
func_config.default_distribute_strategy(flow.scope.consistent_view())
@flow.global_function(func_config)
def eval_generator(input=flow.FixedTensorDef((self.batch_size, 3, 256, 256)),):
def eval_generator(input=flow.FixedTensorDef((self.batch_size, 3, 256, 256))):
g_out = self.generator(input, trainable=False)
return g_out
check_point = flow.train.CheckPoint()
check_point.init()
x, y = self.load_facades()
batch_num = len(x) // self.batch_size
print(batch_num)
for epoch_idx in range(epochs):
for batch_idx in range(batch_num):
label1 = np.ones((self.batch_size, 1, 30, 30)).astype(np.float32)
label0 = np.zeros((self.batch_size, 1, 30, 30)).astype(np.float32)
x, y = self.load_facades()
for batch_idx in range(len(x) // self.batch_size):
inp = x[
batch_idx * self.batch_size : (batch_idx + 1) * self.batch_size
].astype(np.float32, order="C")
target = y[
batch_idx * self.batch_size : (batch_idx + 1) * self.batch_size
].astype(np.float32, order="C")
d_loss = train_discriminator(inp, target, label0, label1).get()
g_loss, _ = train_generator(inp, target, label1).get()
d_loss = train_discriminator(inp, target).get()
g_loss, _ = train_generator(inp, target).get()
if (batch_idx + 1) % self.eval_interval == 0:
print(
"{}th epoch, {}th batch, dloss:{:>12.6f}, gloss:{:>12.6f}".format(
......@@ -449,15 +460,12 @@ if __name__ == "__main__":
os.environ["ENABLE_USER_OP"] = "True"
import argparse
parser = argparse.ArgumentParser(description="flags for multi-node and resource")
parser.add_argument("-g", "--gpu_num_per_node", type=int, default=4, required=False)
parser.add_argument("-g", "--gpu_num_per_node", type=int, default=1, required=False)
parser.add_argument("-e", "--epoch_num", type=int, default=10, required=False)
parser.add_argument("-lr", "--learning_rate", type=float, default=1e-4, required=False)
parser.add_argument("-lr", "--learning_rate", type=float, default=2e-4, required=False)
parser.add_argument(
"-c", "--compare", default=False, action="store_true", required=False
)
parser.add_argument(
"-m", "--multinode", default=False, action="store_true", required=False
)
parser.add_argument(
"-load", "--model_load_dir", type=str, default="checkpoint", required=False
)
......@@ -467,6 +475,5 @@ if __name__ == "__main__":
parser.add_argument("-b", "--batch_size", type=int, default=1, required=False)
args = parser.parse_args()
pix2pix = Pix2Pix(args)
if args.compare:
pix2pix.compare_with_tf()
pix2pix.train(epochs=args.epoch_num)
pix2pix.compare_with_tf()
# pix2pix.train(epochs=args.epoch_num)
Generative/tf_dcgan.py 已删除 100644 → 0
浏览文件 @ 23e573b8
import tensorflow as tf
import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
from tensorflow.keras import layers
import time
from IPython import display
BATCH_SIZE = 32
noise_dim = 100
def const_initializer():
return tf.constant_initializer(0.002)
def make_generator_model():
model = tf.keras.Sequential()
model.add(
layers.Dense(
7 * 7 * 256,
use_bias=False,
input_shape=(100,),
kernel_initializer=const_initializer(),
)
)
model.add(layers.BatchNormalization())
model.add(layers.LeakyReLU())
model.add(layers.Reshape((7, 7, 256)))
assert model.output_shape == (None, 7, 7, 256) # Note: None is the batch size
model.add(
layers.Conv2DTranspose(
128,
(5, 5),
strides=(1, 1),
padding="same",
use_bias=False,
kernel_initializer=const_initializer(),
)
)
assert model.output_shape == (None, 7, 7, 128)
model.add(layers.BatchNormalization())
model.add(layers.LeakyReLU())
model.add(
layers.Conv2DTranspose(
64,
(5, 5),
strides=(2, 2),
padding="same",
use_bias=False,
kernel_initializer=const_initializer(),
)
)
assert model.output_shape == (None, 14, 14, 64)
model.add(layers.BatchNormalization())
model.add(layers.LeakyReLU())
model.add(
layers.Conv2DTranspose(
1,
(5, 5),
strides=(2, 2),
padding="same",
use_bias=False,
kernel_initializer=const_initializer(),
activation="tanh",
)
)
assert model.output_shape == (None, 28, 28, 1)
return model
def make_discriminator_model():
model = tf.keras.Sequential()
model.add(
layers.Conv2D(
64,
(5, 5),
strides=(2, 2),
padding="same",
kernel_initializer=const_initializer(),
input_shape=[28, 28, 1],
)
)
model.add(layers.LeakyReLU())
model.add(layers.Dropout(0.3))
model.add(
layers.Conv2D(
128,
(5, 5),
strides=(2, 2),
padding="same",
kernel_initializer=const_initializer(),
)
)
model.add(layers.LeakyReLU())
model.add(layers.Dropout(0.3))
model.add(layers.Flatten())
model.add(layers.Dense(1, kernel_initializer=const_initializer()))
return model
# cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)
def cross_entropy(x, y):
return tf.nn.sigmoid_cross_entropy_with_logits(logits=y, labels=x)
def generator_loss(fake_output):
return cross_entropy(tf.ones_like(fake_output), fake_output)
def tf_dcgan_test():
generator = make_generator_model()
discriminator = make_discriminator_model()
z = tf.random.normal([BATCH_SIZE, noise_dim])
np.save("z.npy", z)
img = generator(z, training=True)
logit = discriminator(img, training=True)
out = generator_loss(logit)
np.save("out.npy", out.numpy())
Generative/tf_pix2pix.py
浏览文件 @ 8b9075ab
......@@ -6,50 +6,33 @@ import numpy as np
# global vars
OUTPUT_CHANNELS = 3
LAMBDA = 100
BATCH_SIZE = 3
BATCH_SIZE = 1
# download dataset
def get_constant_initializer(constant_value=0.002):
return tf.constant_initializer(constant_value)
def download():
# the default download path is "~/.keras/datasets"
_URL = "https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz"
path_to_zip = tf.keras.utils.get_file("facades.tar.gz", origin=_URL, extract=True)
path_to_zip = tf.keras.utils.get_file(
"facades.tar.gz", origin=_URL, extract=True)
return path_to_zip
# load data, which is a single image with shape (256, 256, 3)
def load(image_file):
image = tf.io.read_file(image_file)
image = tf.image.decode_jpeg(image)
w = tf.shape(image)[1]
w = w // 2
real_image = image[:, :w, :]
input_image = image[:, w:, :]
input_image = tf.cast(input_image, tf.float32)
real_image = tf.cast(real_image, tf.float32)
return input_image, real_image
# build the model
def downsample(inp, filters, size, apply_batchnorm=True, const_init=True):
if not const_init:
initializer = tf.random_normal_initializer(0.0, 0.02)
else:
initializer = tf.constant_initializer(0.002)
conv = tf.keras.layers.Conv2D(
filters,
size,
strides=2,
padding="same",
kernel_initializer=initializer,
kernel_initializer=get_constant_initializer(),
use_bias=False,
)(inp)
if apply_batchnorm:
if apply_batchnorm: #and not const_init:
conv = tf.keras.layers.BatchNormalization()(conv)
result = tf.keras.layers.LeakyReLU()(conv)
......@@ -58,20 +41,16 @@ def downsample(inp, filters, size, apply_batchnorm=True, const_init=True):
def upsample(inp, filters, size, apply_dropout=False, const_init=True):
if not const_init:
initializer = tf.random_normal_initializer(0.0, 0.02)
else:
initializer = tf.constant_initializer(0.002)
deconv = tf.keras.layers.Conv2DTranspose(
filters,
size,
strides=2,
padding="same",
kernel_initializer=initializer,
kernel_initializer=get_constant_initializer(),
use_bias=False,
)(inp)
# if not const_init:
deconv = tf.keras.layers.BatchNormalization()(deconv)
if apply_dropout:
......@@ -83,7 +62,7 @@ def upsample(inp, filters, size, apply_dropout=False, const_init=True):
def Generator(const_init=True):
apply_dropout = False if const_init else True
apply_dropout = False
inputs = tf.keras.layers.Input(shape=[256, 256, 3])
......@@ -119,29 +98,20 @@ def Generator(const_init=True):
u1 = upsample(u2, 64, 4, const_init=const_init) # (bs, 128, 128, 128)
u1 = tf.keras.layers.Concatenate()([u1, d1])
if not const_init:
initializer = tf.random_normal_initializer(0.0, 0.02)
else:
initializer = tf.constant_initializer(0.002)
last = tf.keras.layers.Conv2DTranspose(
OUTPUT_CHANNELS,
4,
strides=2,
padding="same",
kernel_initializer=initializer,
kernel_initializer=get_constant_initializer(),
activation="tanh",
) # (bs, 256, 256, 3)
u0 = last(u1)
return tf.keras.Model(inputs=inputs, outputs=u0)
return tf.keras.Model(inputs=inputs, outputs=d8)
def Discriminator(const_init=True):
if not const_init:
initializer = tf.random_normal_initializer(0.0, 0.02)
else:
initializer = tf.constant_initializer(0.002)
inp = tf.keras.layers.Input(shape=[256, 256, 3], name="input_image")
tar = tf.keras.layers.Input(shape=[256, 256, 3], name="target_image")
......@@ -150,23 +120,25 @@ def Discriminator(const_init=True):
down1 = downsample(
x, 64, 4, apply_batchnorm=False, const_init=const_init
) # (bs, 128, 128, 64)
down2 = downsample(down1, 128, 4, const_init=const_init) # (bs, 64, 64, 128)
down3 = downsample(down2, 256, 4, const_init=const_init) # (bs, 32, 32, 256)
# (bs, 64, 64, 128)
down2 = downsample(down1, 128, 4, const_init=const_init)
# (bs, 32, 32, 256)
down3 = downsample(down2, 256, 4, const_init=const_init)
zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256)
conv = tf.keras.layers.Conv2D(
512, 4, strides=1, kernel_initializer=initializer, use_bias=False
512, 4, strides=1, kernel_initializer=get_constant_initializer(), use_bias=False
)(
zero_pad1
) # (bs, 31, 31, 512)
batchnorm1 = tf.keras.layers.BatchNormalization()(conv, training=True)
bn = tf.keras.layers.BatchNormalization()(conv, training=True)
leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)
leaky_relu = tf.keras.layers.LeakyReLU()(bn)
zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512)
last = tf.keras.layers.Conv2D(1, 4, strides=1, kernel_initializer=initializer)(
last = tf.keras.layers.Conv2D(1, 4, strides=1, kernel_initializer=get_constant_initializer())(
zero_pad2
) # (bs, 30, 30, 1)
......@@ -174,8 +146,9 @@ def Discriminator(const_init=True):
def generator_loss(disc_generated_output, gen_output, target):
loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True)
gan_loss = loss_object(tf.ones_like(disc_generated_output), disc_generated_output)
loss_object = tf.nn.sigmoid_cross_entropy_with_logits
gan_loss = loss_object(tf.ones_like(
disc_generated_output), disc_generated_output)
# mean absolute error
l1_loss = tf.reduce_mean(tf.abs(target - gen_output))
......@@ -186,7 +159,7 @@ def generator_loss(disc_generated_output, gen_output, target):
def discriminator_loss(disc_real_output, disc_generated_output):
loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True)
loss_object = tf.nn.sigmoid_cross_entropy_with_logits
real_loss = loss_object(tf.ones_like(disc_real_output), disc_real_output)
generated_loss = loss_object(
......@@ -198,16 +171,54 @@ def discriminator_loss(disc_real_output, disc_generated_output):
return total_disc_loss
def tf_pix2pix_test():
generator = Generator()
discriminator = Discriminator()
# generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
# discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
generator_optimizer = tf.keras.optimizers.SGD(1e-4)
discriminator_optimizer = tf.keras.optimizers.SGD(1e-4)
@tf.function
def train_step(input_image, target):
# with tf.GradientTape() as disc_tape:
# gen_output = generator(input_image, training=True)
# disc_real_output = discriminator([input_image, target], training=True)
# disc_generated_output = discriminator([input_image, gen_output], training=True)
# disc_loss = discriminator_loss(disc_real_output, disc_generated_output)
# discriminator_gradients = disc_tape.gradient(disc_loss,
# discriminator.trainable_variables)
# discriminator_optimizer.apply_gradients(zip(discriminator_gradients,
# discriminator.trainable_variables))
with tf.GradientTape() as gen_tape:
gen_output = generator(input_image, training=True)
# disc_generated_output = discriminator(
# [input_image, gen_output], training=True)
# gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(
# disc_generated_output, gen_output, target)
generator_gradients = gen_tape.gradient(gen_output,
generator.trainable_variables)
generator_optimizer.apply_gradients(zip(generator_gradients,
generator.trainable_variables))
return gen_output, None
def generate_results():
inp = tf.random.normal([BATCH_SIZE, 256, 256, 3])
tar = tf.random.normal([BATCH_SIZE, 256, 256, 3])
np.save("input.npy", inp)
np.save("target.npy", tar)
generator = Generator()
result = generator(inp, training=True)
discriminator = Discriminator()
result = discriminator([result, tar], training=True)
print(result.shape)
np.save("result.npy", result.numpy())
for i in range(3):
g_loss, d_loss = train_step(inp, tar)
# print("tf produce d_loss:{}, g_loss:{}".format(d_loss.numpy().mean(), g_loss.numpy().mean()))
print("tf produce g_loss:{}".format(g_loss.numpy().mean()))
# np.save("d_loss.npy", d_loss.numpy())
# np.save("g_loss.npy", g_loss.numpy())
if __name__ == '__main__':
generate_results()
Markdown is supported
0% 或 .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑！
想要评论请 注册 或