| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # |
| |
| from singa import device |
| from singa import opt |
| from singa import tensor |
| |
| import argparse |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import os |
| from model import lsgan_mlp |
| from utils import load_data |
| from utils import print_log |
| |
| |
| class LSGAN(): |
| |
| def __init__(self, |
| dev, |
| rows=28, |
| cols=28, |
| channels=1, |
| noise_size=100, |
| hidden_size=128, |
| batch=128, |
| interval=1000, |
| learning_rate=0.001, |
| iterations=1000000, |
| d_steps=3, |
| g_steps=1, |
| dataset_filepath='mnist.pkl.gz', |
| file_dir='lsgan_images/'): |
| self.dev = dev |
| self.rows = rows |
| self.cols = cols |
| self.channels = channels |
| self.feature_size = self.rows * self.cols * self.channels |
| self.noise_size = noise_size |
| self.hidden_size = hidden_size |
| self.batch = batch |
| self.batch_size = self.batch // 2 |
| self.interval = interval |
| self.learning_rate = learning_rate |
| self.iterations = iterations |
| self.d_steps = d_steps |
| self.g_steps = g_steps |
| self.dataset_filepath = dataset_filepath |
| self.file_dir = file_dir |
| self.model = lsgan_mlp.create_model(noise_size=self.noise_size, |
| feature_size=self.feature_size, |
| hidden_size=self.hidden_size) |
| |
| def train(self): |
| train_data, _, _, _, _, _ = load_data(self.dataset_filepath) |
| dev = device.create_cuda_gpu_on(0) |
| dev.SetRandSeed(0) |
| np.random.seed(0) |
| |
| #sgd = opt.SGD(lr=self.learning_rate, momentum=0.9, weight_decay=1e-5) |
| sgd = opt.Adam(lr=self.learning_rate) |
| |
| noise = tensor.Tensor((self.batch_size, self.noise_size), dev, |
| tensor.float32) |
| real_images = tensor.Tensor((self.batch_size, self.feature_size), dev, |
| tensor.float32) |
| real_labels = tensor.Tensor((self.batch_size, 1), dev, tensor.float32) |
| fake_labels = tensor.Tensor((self.batch_size, 1), dev, tensor.float32) |
| substrahend_labels = tensor.Tensor((self.batch_size, 1), dev, tensor.float32) |
| |
| # attached model to graph |
| self.model.set_optimizer(sgd) |
| self.model.compile([noise], |
| is_train=True, |
| use_graph=False, |
| sequential=True) |
| |
| real_labels.set_value(1.0) |
| fake_labels.set_value(-1.0) |
| substrahend_labels.set_value(0.0) |
| |
| for iteration in range(self.iterations): |
| |
| for d_step in range(self.d_steps): |
| idx = np.random.randint(0, train_data.shape[0], self.batch_size) |
| real_images.copy_from_numpy(train_data[idx]) |
| |
| self.model.train() |
| |
| # Training the Discriminative Net |
| _, d_loss_real = self.model.train_one_batch_dis( |
| real_images, real_labels) |
| |
| noise.uniform(-1, 1) |
| fake_images = self.model.forward_gen(noise) |
| _, d_loss_fake = self.model.train_one_batch_dis( |
| fake_images, fake_labels) |
| |
| d_loss = tensor.to_numpy(d_loss_real)[0] + tensor.to_numpy( |
| d_loss_fake)[0] |
| |
| for g_step in range(self.g_steps): |
| # Training the Generative Net |
| noise.uniform(-1, 1) |
| _, g_loss_tensor = self.model.train_one_batch( |
| noise, substrahend_labels) |
| |
| g_loss = tensor.to_numpy(g_loss_tensor)[0] |
| |
| if iteration % self.interval == 0: |
| self.model.eval() |
| self.save_image(iteration) |
| print_log(' The {} iteration, G_LOSS: {}, D_LOSS: {}'.format( |
| iteration, g_loss, d_loss)) |
| |
| def save_image(self, iteration): |
| demo_row = 5 |
| demo_col = 5 |
| if not hasattr(self, "demo_noise"): |
| self.demo_noise = tensor.Tensor( |
| (demo_col * demo_row, self.noise_size), dev, tensor.float32) |
| self.demo_noise.uniform(-1, 1) |
| gen_imgs = self.model.forward_gen(self.demo_noise) |
| gen_imgs = tensor.to_numpy(gen_imgs) |
| show_imgs = np.reshape( |
| gen_imgs, (gen_imgs.shape[0], self.rows, self.cols, self.channels)) |
| fig, axs = plt.subplots(demo_row, demo_col) |
| cnt = 0 |
| for r in range(demo_row): |
| for c in range(demo_col): |
| axs[r, c].imshow(show_imgs[cnt, :, :, 0], cmap='gray') |
| axs[r, c].axis('off') |
| cnt += 1 |
| fig.savefig("{}{}.png".format(self.file_dir, iteration)) |
| plt.close() |
| |
| |
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser(description='Train GAN over MNIST') |
| parser.add_argument('filepath', type=str, help='the dataset path') |
| parser.add_argument('--use_gpu', action='store_true') |
| args = parser.parse_args() |
| |
| if args.use_gpu: |
| print('Using GPU') |
| dev = device.create_cuda_gpu() |
| else: |
| print('Using CPU') |
| dev = device.get_default_device() |
| |
| if not os.path.exists('lsgan_images/'): |
| os.makedirs('lsgan_images/') |
| |
| rows = 28 |
| cols = 28 |
| channels = 1 |
| noise_size = 100 |
| hidden_size = 128 |
| batch = 128 |
| interval = 1000 |
| learning_rate = 0.0005 |
| iterations = 1000000 |
| d_steps = 1 |
| g_steps = 1 |
| dataset_filepath = 'mnist.pkl.gz' |
| file_dir = 'lsgan_images/' |
| lsgan = LSGAN(dev, rows, cols, channels, noise_size, hidden_size, batch, |
| interval, learning_rate, iterations, d_steps, g_steps, |
| dataset_filepath, file_dir) |
| lsgan.train() |
