def main():
import numpy as np
import argparse
# パーサーを作る
parser = argparse.ArgumentParser(
prog='train', # プログラム名
usage='train DCGAN to dish', # プログラムの利用方法
description='description', # 引数のヘルプの前に表示
epilog='end', # 引数のヘルプの後で表示
add_help=True, # -h/–help オプションの追加
)
# 引数の追加
parser.add_argument('-s', '--seed', help='seed', type=int, required=True)
parser.add_argument('-n',
'--number',
help='the number of experiments.',
type=int,
required=True)
parser.add_argument('--hidden',
help='the number of codes of Generator.',
type=int,
default=100)
parser.add_argument('-e',
'--epoch',
help='the number of epoch, defalut value is 300',
type=int,
default=300)
parser.add_argument('-bs',
'--batch_size',
help='batch size. defalut value is 128',
type=int,
default=128)
parser.add_argument('-g',
'--gpu',
help='specify gpu by this number. defalut value is 0',
choices=[0, 1],
type=int,
default=0)
parser.add_argument(
'-ks',
'--ksize',
help='specify ksize of generator by this number. any of following;'
' 4 or 6. defalut value is 6',
choices=[4, 6],
type=int,
default=6)
parser.add_argument(
'-dis',
'--discriminator',
help='specify discriminator by this number. any of following;'
' 0: original, 1: minibatch discriminatio, 2: feature matching, 3: GAP. defalut value is 3',
choices=[0, 1, 2, 3],
type=int,
default=3)
parser.add_argument(
'-ts',
'--tensor_shape',
help=
'specify Tensor shape by this numbers. first args denotes to B, seconds to C.'
' defalut value are B:32, C:8',
type=int,
default=[32, 8],
nargs=2)
parser.add_argument('-V',
'--version',
version='%(prog)s 1.0.0',
action='version',
default=False)
# 引数を解析する
args = parser.parse_args()
gpu = args.gpu
batch_size = args.batch_size
n_hidden = args.hidden
epoch = args.epoch
seed = args.seed
number = args.number # number of experiments
if args.ksize == 6:
pad = 2
else:
pad = 1
out = pathlib.Path("result_{0}".format(number))
if not out.exists():
out.mkdir()
out /= pathlib.Path("result_{0}_{1}".format(number, seed))
if not out.exists():
out.mkdir()
# 引数(ハイパーパラメータの設定)の書き出し
with open(out / "args.txt", "w") as f:
f.write(str(args))
print('GPU: {}'.format(gpu))
print('# Minibatch-size: {}'.format(batch_size))
print('# n_hidden: {}'.format(n_hidden))
print('# epoch: {}'.format(epoch))
print('# out: {}'.format(out))
print('# seed: {}'.format(seed))
print('# ksize: {}'.format(args.ksize))
print('# pad: {}'.format(pad))
# fix seed
np.random.seed(seed)
if chainer.backends.cuda.available:
chainer.backends.cuda.cupy.random.seed(seed)
# import discrimination & set up
# if args.discriminator == 0:
# print("# Original Discriminator")
# from discriminator import Discriminator
# from updater import DCGANUpdater
# dis = Discriminator()
if args.discriminator == 1:
print("# Discriminator applied Minibatch Discrimination")
print('# Tensor shape is A x {0} x {1}'.format(args.tensor_shape[0],
args.tensor_shape[1]))
from discriminator_md import Discriminator
from updater import DCGANUpdater
dis = Discriminator(B=args.tensor_shape[0], C=args.tensor_shape[1])
elif args.discriminator == 3:
print("# Discriminator applied GAP")
from discriminator import Discriminator
from updater import DCGANUpdater
dis = Discriminator()
"""
elif args.discriminator == 2:
print("# Discriminator applied matching")
from discriminator_fm import Discriminator
from updater_fm import DCGANUpdater
"""
print('')
# Set up a neural network to train
gen = Generator(n_hidden=n_hidden, ksize=args.ksize, pad=pad)
if gpu >= 0:
# Make a specified GPU current
chainer.backends.cuda.get_device_from_id(gpu).use()
gen.to_gpu() # Copy the model to the GPU
dis.to_gpu()
opt_gen = make_optimizer(gen)
opt_dis = make_optimizer(dis)
# Prepare Dataset
"""
paths = ["rsize_data_128", "test_rsize_data_128",
"unlabeled_rsize_data_128"] # resize data 128
"""
paths = ["center_crop_data_128"] # center ctop data
data_path = []
for path in paths:
data_dir = pathlib.Path(path)
abs_data_dir = data_dir.resolve()
print("data dir path:", abs_data_dir)
data_path += [path for path in abs_data_dir.glob("*.jpg")]
print("data length:", len(data_path))
data = ImageDataset(paths=data_path) # dtype=np.float32
train_iter = chainer.iterators.SerialIterator(data, batch_size)
# Set up a updater and trainer
updater = DCGANUpdater(models=(gen, dis),
iterator=train_iter,
optimizer={
'gen': opt_gen,
'dis': opt_dis
},
device=gpu)
trainer = training.Trainer(updater, (epoch, 'epoch'), out=out)
snapshot_interval = (10, 'epoch')
display_interval = (1, 'epoch')
# storage method is hdf5
trainer.extend(extensions.snapshot(
filename='snapshot_epoch_{.updater.epoch}.npz', savefun=save_npz),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(gen,
'gen_epoch_{.updater.epoch}.npz',
savefun=save_npz),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(dis,
'dis_epoch_{.updater.epoch}.npz',
savefun=save_npz),
trigger=snapshot_interval)
trainer.extend(extensions.LogReport())
trainer.extend(extensions.PrintReport([
'epoch', 'iteration', 'gen/loss', 'dis/loss', 'elapsed_time',
'dis/accuracy'
]),
trigger=display_interval)
trainer.extend(extensions.ProgressBar(update_interval=20))
trainer.extend(out_generated_image(gen, dis, 5, 5, seed, out),
trigger=display_interval)
trainer.extend(
extensions.PlotReport(['gen/loss', 'dis/loss'],
x_key='epoch',
file_name='loss_{0}_{1}.jpg'.format(
number, seed),
grid=False))
trainer.extend(extensions.dump_graph("gen/loss", out_name="gen.dot"))
trainer.extend(extensions.dump_graph("dis/loss", out_name="dis.dot"))
# Run the training
trainer.run()
print("# Discriminator applied GAP")
from discriminator_gap import Discriminator
from updater import DCGANUpdater
dis = Discriminator()
print('')
# Set up a neural network to train
gen = Generator(n_hidden=n_hidden)
dis = Discriminator()
if gpu >= 0:
# Make a specified GPU current
chainer.backends.cuda.get_device_from_id(gpu).use()
gen.to_gpu() # Copy the model to the GPU
dis.to_gpu()
# Setup an optimizer
def make_optimizer(model, alpha=0.0002, beta1=0.5):
optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
optimizer.setup(model)
# optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec')
return optimizer
opt_gen = make_optimizer(gen)
opt_dis = make_optimizer(dis)
# Load the GMM dataset
dataset, _ = chainer.datasets.get_svhn(withlabel=False, scale=255.)
train_iter = chainer.iterators.SerialIterator(dataset, batch_size)
print("# Data size: {}".format(len(dataset)), end="\n\n")
def main():
parser = argparse.ArgumentParser(description='Chainer: DCGAN MNIST')
parser.add_argument('--batchsize',
'-b',
type=int,
default=50,
help='Number of images in each mini-batch')
parser.add_argument('--epoch',
'-e',
type=int,
default=1000,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu',
'-g',
type=int,
default=0,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--out',
'-o',
default='result',
help='Directory to output the result')
parser.add_argument('--resume',
'-r',
default='',
help='Resume the training from snapshot')
parser.add_argument('--n_hidden',
'-n',
type=int,
default=100,
help='Number of hidden units (z)')
parser.add_argument('--seed',
type=int,
default=0,
help='Random seed of z at visualization stage')
parser.add_argument('--snapshot_interval',
type=int,
default=1000,
help='Interval of snapshot')
parser.add_argument('--display_interval',
type=int,
default=100,
help='Interval of displaying log to console')
parser.add_argument('--image_dir',
default='./img',
help='Training image directory')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# n_hidden: {}'.format(args.n_hidden))
print('# epoch: {}'.format(args.epoch))
print('')
gen = Generator()
dis = Discriminator()
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
gen.to_gpu()
dis.to_gpu()
def make_optimizer(model, alpha=0.0002, beta1=0.5):
optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec')
return optimizer
opt_gen = make_optimizer(gen)
opt_dis = make_optimizer(dis)
def read_dataset(image_dir):
fs = os.listdir(image_dir)
dataset = []
for fn in fs:
img = Image.open(f"{args.image_dir}/{fn}").convert('RGB')
dataset.append((np.asarray(img,
dtype=np.float32).transpose(2, 0, 1)))
return dataset
dataset = read_dataset(args.image_dir)
train_iter = chainer.iterators.SerialIterator(dataset,
args.batchsize,
repeat=True,
shuffle=True)
updater = DCGANUpdater(models=(gen, dis),
iterator=train_iter,
optimizer={
'gen': opt_gen,
'dis': opt_dis
},
device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
epoch_interval = (1, 'epoch')
# snapshot_interval = (args.snapshot_interval, 'iteration')
display_interval = (args.display_interval, 'iteration')
# trainer.extend(extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
# trainer.extend(extensions.snapshot_object(gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
# trainer.extend(extensions.snapshot_object(dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
trainer.extend(
extensions.snapshot(filename='snapshot_epoch_{.updater.epoch}.npz'),
trigger=epoch_interval)
trainer.extend(extensions.snapshot_object(
gen, 'gen_epoch_{.updater.epoch}.npz'),
trigger=epoch_interval)
trainer.extend(extensions.snapshot_object(
dis, 'dis_epoch_{.updater.epoch}.npz'),
trigger=epoch_interval)
trainer.extend(extensions.LogReport(trigger=display_interval))
trainer.extend(extensions.PrintReport([
'epoch',
'iteration',
'gen/loss',
'dis/loss',
]),
trigger=display_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(out_generated_image(gen, dis, 10, 10, args.seed, args.out),
trigger=epoch_interval)
if args.resume:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
