def main():
parser = argparse.ArgumentParser(description='chainer implementation of pix2pix')
parser.add_argument('--gpu', '-g', type=int, default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--seed', type=int, default=0,
help='Random seed')
parser.add_argument('--model', '-m', default='',
help='model snapshot')
parser.add_argument('--enc', '-e', type=str, default='enc_iter_60000.npz', help='encoder snapshot')
parser.add_argument('--dec', '-d', type=str, default='dec_iter_60000.npz', help='decoder snapshot')
parser.add_argument('--out', '-o', type=str, default='out', help='output dir')
parser.add_argument('--input', '-i', default='sample.jpg', help='input jpg', required=True)
parser.add_argument('--contour', '-c', action='store_true', help='from contour image or not')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
enc = Encoder(in_ch=3)
dec = Decoder(out_ch=3)
dis = Discriminator(in_ch=3, out_ch=3)
if args.gpu >= 0:
chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current
enc.to_gpu() # Copy the model to the GPU
dec.to_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.00001), 'hook_dec')
return optimizer
if args.model:
opt_enc = make_optimizer(enc)
opt_dec = make_optimizer(dec)
opt_dis = make_optimizer(dis)
# Set up a trainer
updater = FacadeUpdater(
models=(enc, dec, dis),
iterator={},
optimizer={
'enc': opt_enc, 'dec': opt_dec,
'dis': opt_dis},
device=args.gpu)
trainer = training.Trainer(updater, (200, 'epoch'), out='generate/')
chainer.serializers.load_npz(args.model, trainer)
elif args.enc and args.dec:
chainer.serializers.load_npz(args.enc, enc)
chainer.serializers.load_npz(args.dec, dec)
if not args.contour:
from make_contour import get_contour_image
get_contour_image(args.input)
generate_image_from_contour(args.input, enc, dec, args.out)
def discriminate():
dis = Discriminator()
chainer.backends.cuda.get_device_from_id(0).use()
serializers.load_npz("./result/dis_iter_25000.npz", dis)
dis.to_gpu() # Copy the model to the GPU
paths = ["./data/70.png", "./data/35.png"]
for path in paths:
img = read_image_as_array(path, dtype=np.float32)
img = img.reshape(-1, img.shape[0], img.shape[1],
img.shape[2]).transpose(0, 3, 1, 2)
img = Variable(chainer.backends.cuda.to_gpu(img))
with chainer.using_config('train', False):
x = dis(img)
x = chainer.backends.cuda.to_cpu(x.data)
print(x)
def train(args):
# setting for logging
if not os.path.exists(args.log):
os.mkdir(args.log)
logger = logging.getLogger()
logging.basicConfig(level=logging.INFO)
log_path = os.path.join(args.log, 'log')
file_handler = logging.FileHandler(log_path)
fmt = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
file_handler.setFormatter(fmt)
logger.addHandler(file_handler)
logger.info('Arguments...')
for arg, val in vars(args).items():
logger.info('{:>10} -----> {}'.format(arg, val))
x, y = gen_synthetic_data(DIM, DIM_EMB, NUM)
train_x, test_x, train_y, test_y = train_test_split(x, y, test_size=0.2)
valid_x, test_x, valid_y, test_y = train_test_split(test_x, test_y, test_size=0.5)
gen = Generator(DIM_EMB)
dis = Discriminator(DIM_EMB)
gen_opt = optimizers.Adam()
dis_opt = optimizers.Adam()
gen_opt.setup(gen)
dis_opt.setup(dis)
trainer = GANTrainer((gen, dis), (gen_opt, dis_opt), logger, (valid_x, valid_y), args.epoch)
trainer.fit(train_x, train_y)
def run_seq_gan():
config = Config()
n_samples = config.get('n_samples')
batch_size = config.get('batch_size')
gen_embedding_dim = config.get('gen_embedding_dim')
gen_hidden_dim = config.get('gen_hidden_dim')
dis_embedding_dim = config.get('dis_embedding_dim')
dis_hidden_dim = config.get('dis_hidden_dim')
dataset_features = config.get('dataset_features')
dataset_dtypes = config.get('dataset_dtypes')
generated_features = config.get('generated_features')
service_list = config.get('service_list')
protocol_service_dict = config.get('protocol_service_dict')
service_port_dict = config.get('service_port_dict')
file_path = config.get('file_path')
CUDA = torch.cuda.is_available()
dataset = Traffic_Dataset(file_path,
dataset_features,
dataset_dtypes,
generated_features,
batch_size=batch_size,
transform=build_input_indices)
vocab_dim = dataset.vocabulary_length
max_seq_len = dataset.max_seq_length
train_epochs = 100
g = Generator(gen_embedding_dim, gen_hidden_dim, vocab_dim, max_seq_len,
CUDA)
d = Discriminator(dis_embedding_dim, dis_hidden_dim, vocab_dim,
max_seq_len, CUDA)
if CUDA:
g.cuda()
d.cuda()
g_opt = optim.Adam(g.parameters())
d_opt = optim.Adagrad(d.parameters())
pre_training(g, d, g_opt, d_opt, dataset, n_samples, batch_size, CUDA)
training(g, d, g_opt, d_opt, dataset, train_epochs, n_samples, batch_size,
CUDA, service_list, protocol_service_dict, service_port_dict)
visualize(dataset_features, generated_features)
def __init__(self, opt):
self.opt = opt
self.global_step = opt.load_iter
self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# define net class instance
self.G_net = Generator(opt).to(self.device)
self.D_net = Discriminator(opt).to(self.device)
if opt.load_model and opt.load_iter > 0:
self._load_pre_model(self.G_net, 'G')
self._load_pre_model(self.D_net, 'D')
# define objectives and optimizers
self.adv_loss = torch.mean # 这里的adv loss直接是真假结果的评分，真图越大越好，假图越小越好
# self.cls_loss = torch.nn.BCELoss() # ??????????? 有啥区别
self.cls_loss = F.binary_cross_entropy_with_logits
self.rec_loss = torch.mean
self.G_optimizer = torch.optim.Adam(self.G_net.parameters(), opt.G_lr, [opt.beta1, opt.beta2])
self.D_optimizer = torch.optim.Adam(self.D_net.parameters(), opt.D_lr, [opt.beta1, opt.beta2])
self.sample_gotten = False # 把它放在init里面，是因为它只随着类的调用初始化一次，是固定的sample
self.writer = TBVisualizer(opt)
def test_discriminate_fakevideo(self):
dis = Discriminator()
gen = Generator()
z = Variable(np.asarray(gen.make_noize(self.batchsize)))
self.assertEqual((self.batchsize, 100), z.shape)
x_fake = gen(z)
self.assertEqual(
(self.batchsize, 3, self.frame, self.height, self.width),
x_fake.shape)
y_fake = dis(x_fake)
self.assertEqual((self.batchsize, 1), y_fake.shape)
def __init__(self, args):
self.args = args
self.writer = SummaryWriter(args.output_dir)
self.iter_i = 1
# data
transform_list = [
transforms.Resize(args.imsize),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]
train_dataset = datasets.CIFAR10(
'./data',
train=True,
transform=transforms.Compose(transform_list),
download=True)
self.train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
num_workers=args.n_workers,
shuffle=True)
# network
self.G = Generator(args.nz, args.ngf, args.nc).to(args.device)
self.D = Discriminator(args.nc, args.ndf).to(args.device)
self.criterion = nn.BCELoss()
self.optimizer_G = optim.Adam(self.G.parameters(),
lr=args.lr,
betas=(args.beta1, args.beta2))
self.optimizer_D = optim.Adam(self.D.parameters(),
lr=args.lr,
betas=(args.beta1, args.beta2))
self.real_label = 1
self.fake_label = 0
self.fixed_z = torch.randn((args.batch_size, args.nz, 1, 1),
device=args.device)
def setup_discriminator(config):
from net import Discriminator, BigBiGANDiscriminator
num_z = 1 if config.generator_architecture == "dcgan" else 2
if hasattr(config, "bigan") and config.bigan:
discriminator = BigBiGANDiscriminator(config.ch,
config.ch * num_z,
enable_blur=config.enable_blur,
sn=config.sn,
res=config.res_dis)
else:
discriminator = Discriminator(ch=config.ch,
enable_blur=config.enable_blur,
sn=config.sn,
res=config.res_dis)
return discriminator
def test_realvideo(self):
dis = Discriminator()
all_files = os.listdir(self.dataset)
video_files = [f for f in all_files if ('mp4' in f)]
train = PreprocessedDataset(paths=video_files, root=self.dataset)
train_iter = chainer.iterators.SerialIterator(train, self.batchsize)
batch = train_iter.next()
x_real = Variable(convert.concat_examples(batch, self.gpu))
self.assertEqual(
(self.batchsize, 3, self.frame, self.height, self.width),
x_real.shape)
y_real = dis(x_real)
self.assertEqual((self.batchsize, 1), y_real.shape)
def main():
gpu_id = 0
batchsize = 10
report_keys = ["loss_dis", "loss_gen"]
train_dataset = datasets.LabeledImageDataset(str(argv[1]))
train_iter = iterators.SerialIterator(train_dataset, batchsize)
models = []
generator = Generator()
discriminator = Discriminator()
opts = {}
opts["opt_gen"] = setup_adam_optimizer(generator)
opts["opt_dis"] = setup_adam_optimizer(discriminator)
models = [generator, discriminator]
chainer.cuda.get_device_from_id(gpu_id).use()
print("use gpu {}".format(gpu_id))
for m in models:
m.to_gpu()
updater_args = {
"iterator": {
'main': train_iter
},
"device": gpu_id,
"optimizer": opts,
"models": models
}
output = 'result'
display_interval = 20
evaluation_interval = 1000
max_iter = 10000
updater = Updater(**updater_args)
trainer = training.Trainer(updater, (max_iter, 'iteration'), out=output)
trainer.extend(
extensions.LogReport(keys=report_keys,
trigger=(display_interval, 'iteration')))
trainer.extend(extensions.PrintReport(report_keys),
trigger=(display_interval, 'iteration'))
trainer.extend(sample_generate(generator, output),
trigger=(evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.run()
help='input model file path without extension')
parser.add_argument('--output', '-o', required=True, type=str,
help='output model file path without extension')
parser.add_argument('--iter', default=100, type=int,
help='number of iteration')
parser.add_argument('--out_image_dir', default=None, type=str,
help='output directory to output images')
parser.add_argument('--dataset', '-d', default='dataset/etl9g.pkl', type=str,
help='dataset file path')
args = parser.parse_args()
gen_model = Generator()
optimizer_gen = optimizers.Adam(alpha=0.0002, beta1=0.5)
optimizer_gen.setup(gen_model)
optimizer_gen.add_hook(chainer.optimizer.WeightDecay(0.00001))
dis_model = Discriminator()
optimizer_dis = optimizers.Adam(alpha=0.0002, beta1=0.5)
optimizer_dis.setup(dis_model)
optimizer_dis.add_hook(chainer.optimizer.WeightDecay(0.00001))
if args.input != None:
serializers.load_hdf5(args.input + '.gen.model', gen_model)
serializers.load_hdf5(args.input + '.gen.state', optimizer_gen)
serializers.load_hdf5(args.input + '.dis.model', dis_model)
serializers.load_hdf5(args.input + '.dis.state', optimizer_dis)
if args.out_image_dir != None:
if not os.path.exists(args.out_image_dir):
try:
os.mkdir(args.out_image_dir)
except:
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--batchsize', '-b', type=int, default=32,
help='Number of images in each mini-batch')
parser.add_argument('--epoch', '-e', type=int, default=100,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu', '-g', type=int, default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--source', default='mnist')
parser.add_argument('--target', default='mnist_m')
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('--seed', type=int, default=0,
help='Random seed of z at visualization stage')
parser.add_argument('--n_processes', type=int, default=16,
help='Interval of displaying log to console')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# epoch: {}'.format(args.epoch))
print('')
# Set up a neural network to train
gen = Generator(
n_hidden=params['gen']['n_hidden'],
n_resblock=params['gen']['n_resblock'],
n_ch=params['gen']['n_ch'], wscale=params['gaussian_wscale'],
res=params['res_image'], bn_eps=params['bn_eps'])
dis = Discriminator(
n_ch=params['dis']['n_ch'], wscale=params['gaussian_wscale'],
bn_eps=params['bn_eps'], dr_prob=params['dis']['dropout_prob'],
noise_sigma=params['dis']['noise_sigma']
)
cls = L.Classifier(DigitClassifier(n_class=params['n_class']))
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
gen.to_gpu()
dis.to_gpu()
cls.to_gpu()
p_opt = params['optimize']
# Setup an optimizer
def make_optimizer(model):
optimizer = chainer.optimizers.Adam(alpha=p_opt['base_lr'],
beta1=p_opt['beta1'])
optimizer.setup(model)
optimizer.add_hook(
chainer.optimizer.WeightDecay(p_opt['weight_decay']))
return optimizer
opt_gen = make_optimizer(gen)
opt_dis = make_optimizer(dis)
opt_cls = make_optimizer(cls)
def load_dataset(name, dtype='train'):
if name == 'mnist':
train, _ = chainer.datasets.get_mnist(withlabel=True, ndim=3)
dataset = TransformDataset(train, transform=gray2rgb)
return TransformDataset(dataset, transform=scale)
elif name == 'mnist_m':
dataset = get_mnist_m(dtype, withlabel=True)
return TransformDataset(dataset, transform=scale)
else:
raise NotImplementedError
source = load_dataset(args.source)
# from chainer.datasets import split_dataset
# source, _ = split_dataset(source, split_at=1000)
target_train = load_dataset(args.target, dtype='train')
source_iter = MultiprocessIterator(
source, args.batchsize, n_processes=args.n_processes)
target_train_iter = MultiprocessIterator(
target_train, args.batchsize, n_processes=args.n_processes)
# Set up a trainer
updater = PixelDAUpdater(
models=(gen, dis, cls),
iterator={'main': source_iter, 'target': target_train_iter},
optimizer={
'gen': opt_gen, 'dis': opt_dis, 'cls': opt_cls},
device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
snapshot_interval = (1, 'epoch')
display_interval = (10, 'iteration')
for opt in [opt_gen, opt_cls, opt_dis]:
trainer.extend(
extensions.ExponentialShift('alpha', p_opt['alpha_decay_rate'],
optimizer=opt),
trigger=(p_opt['alpha_decay_steps'], 'iteration'))
trainer.extend(
extensions.snapshot(filename='snapshot_epoch_{.updater.epoch}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.LogReport(trigger=display_interval))
trainer.extend(extensions.PrintReport([
'epoch', 'iteration', 'gen/loss', 'dis/loss', 'cls/loss',
'validation/main/accuracy'
]), trigger=display_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
# Save two plot images to the result dir
if extensions.PlotReport.available():
trainer.extend(
extensions.PlotReport(['gen/loss', 'dis/loss', 'cls/loss'],
'iteration', trigger=(100, 'iteration'),
file_name='loss.png'))
trainer.extend(
extensions.PlotReport(['validation/main/accuracy'],
'epoch', file_name='accuracy.png'))
# Dump examples of generated images for every epoch
trainer.extend(out_generated_image(source_iter, gen, args.gpu, args.out))
# Evaluate the model with the test dataset for each epoch
target_test = load_dataset(args.target, dtype='test')
target_test_iter = MultiprocessIterator(
target_test, args.batchsize, n_processes=args.n_processes,
repeat=False, shuffle=False)
trainer.extend(
extensions.Evaluator(target_test_iter, cls, device=args.gpu))
# Visualize computational graph for debug
# trainer.extend(extensions.dump_graph('gen/loss', out_name='gen.dot'))
# trainer.extend(extensions.dump_graph('dis/loss', out_name='dis.dot'))
# trainer.extend(extensions.dump_graph('cls/loss', out_name='cls.dot'))
if args.resume:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
def main():
parser = argparse.ArgumentParser(
description='chainer implementation of pix2pix')
parser.add_argument('--batchsize', '-b', type=int, default=1,
help='Number of images in each mini-batch')
parser.add_argument('--epoch', '-e', type=int, default=200,
help='Number of sweeps over the dataset to train')
parser.add_argument('--device', '-d', type=str, default='-1',
help='Device specifier. Either ChainerX device '
'specifier or an integer. If non-negative integer, '
'CuPy arrays with specified device id are used. If '
'negative integer, NumPy arrays are used')
parser.add_argument('--dataset', '-i', default='./facade/base',
help='Directory of image files.')
parser.add_argument('--out', '-o', default='result',
help='Directory to output the result')
parser.add_argument('--resume', '-r', type=str,
help='Resume the training from snapshot')
parser.add_argument('--seed', type=int, default=0,
help='Random seed')
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')
group = parser.add_argument_group('deprecated arguments')
group.add_argument('--gpu', '-g', dest='device',
type=int, nargs='?', const=0,
help='GPU ID (negative value indicates CPU)')
args = parser.parse_args()
if chainer.get_dtype() == numpy.float16:
warnings.warn(
'This example may cause NaN in FP16 mode.', RuntimeWarning)
device = chainer.get_device(args.device)
if device.xp is chainerx:
sys.stderr.write('This example does not support ChainerX devices.\n')
sys.exit(1)
print('Device: {}'.format(device))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# epoch: {}'.format(args.epoch))
print('')
device.use()
# Set up a neural network to train
enc = Encoder(in_ch=12)
dec = Decoder(out_ch=3)
dis = Discriminator(in_ch=12, out_ch=3)
enc.to_device(device)
dec.to_device(device)
dis.to_device(device)
# 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.00001), 'hook_dec')
return optimizer
opt_enc = make_optimizer(enc)
opt_dec = make_optimizer(dec)
opt_dis = make_optimizer(dis)
train_d = FacadeDataset(args.dataset, data_range=(1, 300))
test_d = FacadeDataset(args.dataset, data_range=(300, 379))
train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize)
test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize)
# Set up a trainer
updater = FacadeUpdater(
models=(enc, dec, dis),
iterator={
'main': train_iter,
'test': test_iter},
optimizer={
'enc': opt_enc, 'dec': opt_dec,
'dis': opt_dis},
device=device)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
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(
enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
trainer.extend(extensions.LogReport(trigger=display_interval))
trainer.extend(extensions.PrintReport([
'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss',
]), trigger=display_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(
out_image(
updater, enc, dec,
5, 5, args.seed, args.out),
trigger=snapshot_interval)
if args.resume is not None:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
ndf = 64
lr = 0.0002
beta1 = 0.5
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
torch.nn.init.normal_(m.weight, 0.0, 0.02)
elif classname.find('BatchNorm') != -1:
torch.nn.init.normal_(m.weight, 1.0, 0.02)
torch.nn.init.zeros_(m.bias)
netG = Generator(ngpu, nz, ngf).to(device)
netG.apply(weights_init)
netD = Discriminator(ngpu, ndf).to(device)
netD.apply(weights_init)
criterion = nn.BCELoss()
fixed_noise = torch.randn(opt.batchSize, nz, 1, 1, device=device)
real_label = 1
fake_label = 0
optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))
schedulerD = ReduceLROnPlateau(optimizerD,
mode='min',
factor=0.5,
patience=2,
verbose=True)
def main():
FLAGS(sys.argv)
running_helper = RunningHelper(FLAGS.use_mpi)
global mpi_is_master
mpi_is_master = running_helper.is_master
# Check stage / image size / dynamic batch size / data consistency.
running_helper.check_hps_consistency()
# Setup Models
mapping = MappingNetwork(FLAGS.ch)
generator = StyleGenerator(FLAGS.ch, enable_blur=FLAGS.enable_blur)
discriminator = Discriminator(ch=FLAGS.ch, enable_blur=FLAGS.enable_blur)
if running_helper.keep_smoothed_gen:
smoothed_generator = StyleGenerator(FLAGS.ch,
enable_blur=FLAGS.enable_blur)
smoothed_mapping = MappingNetwork(FLAGS.ch)
models = [mapping, generator, discriminator]
model_names = ['Mapping', 'Generator', 'Discriminator']
if running_helper.keep_smoothed_gen:
models.append(smoothed_generator)
models.append(smoothed_mapping)
model_names.append('SmoothedGenerator')
model_names.append('SmoothedMapping')
if running_helper.device > -1:
chainer.cuda.get_device_from_id(running_helper.device).use()
for model in models:
model.to_gpu()
stage_manager = StageManager(
stage_interval=running_helper.stage_interval,
dynamic_batch_size=running_helper.dynamic_batch_size,
make_dataset_func=running_helper.make_dataset,
make_iterator_func=make_iterator_func,
debug_start_instance=FLAGS.debug_start_instance)
#if running_helper.is_master:
# chainer.global_config.debug = True
updater_args = {
"models": models,
"optimizer": {
"map":
running_helper.make_optimizer(mapping, FLAGS.adam_alpha_g / 100,
FLAGS.adam_beta1, FLAGS.adam_beta2),
"gen":
running_helper.make_optimizer(generator, FLAGS.adam_alpha_g,
FLAGS.adam_beta1, FLAGS.adam_beta2),
"dis":
running_helper.make_optimizer(discriminator, FLAGS.adam_alpha_d,
FLAGS.adam_beta1, FLAGS.adam_beta2)
},
'stage_manager': stage_manager,
'lambda_gp': FLAGS.lambda_gp,
'smoothing': FLAGS.smoothing,
'style_mixing_rate': FLAGS.style_mixing_rate,
'use_cleargrads': running_helper.use_cleargrads,
'total_gpu': running_helper.fleet_size
}
updater = Updater(**updater_args)
trainer = training.Trainer(updater,
(lambda _trainer: _trainer.updater.stage_manager
.stage_int >= FLAGS.max_stage),
out=FLAGS.out)
# Set up extensions
if running_helper.is_master:
for model, model_name in zip(models, model_names):
trainer.extend(extensions.snapshot_object(
model, model_name + '_{.updater.iteration}.npz'),
trigger=(FLAGS.snapshot_interval, 'iteration'))
trainer.extend(extensions.snapshot(
filename='snapshot_iter_{.updater.iteration}.npz'),
trigger=(FLAGS.snapshot_interval, 'iteration'))
trainer.extend(
extensions.ProgressBar(training_length=(updater.total_iteration,
'iteration'),
update_interval=1))
trainer.extend(sample_generate_light(generator,
mapping,
FLAGS.out,
rows=8,
cols=8),
trigger=(FLAGS.evaluation_sample_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
if running_helper.keep_smoothed_gen:
trainer.extend(sample_generate_light(smoothed_generator,
smoothed_mapping,
FLAGS.out,
rows=8,
cols=8,
subdir='preview_smoothed'),
trigger=(FLAGS.evaluation_sample_interval,
'iteration'),
priority=extension.PRIORITY_WRITER)
report_keys = [
'iteration', 'elapsed_time', 'stage', 'batch_size', 'image_size',
'gen/loss_adv', 'dis/loss_adv', 'dis/loss_gp'
]
if FLAGS.fid_interval > 0:
report_keys += 'FID'
fidapi = FIDAPI(FLAGS.fid_clfs_type,
FLAGS.fid_clfs_path,
gpu=running_helper.device,
load_real_stat=FLAGS.fid_real_stat)
trainer.extend(fid_extension(fidapi,
batch_generate_func(
generator, mapping, trainer),
seed=FLAGS.seed,
report_key='FID'),
trigger=(FLAGS.fid_interval, 'iteration'))
if running_helper.keep_smoothed_gen:
report_keys += 'S_FID'
trainer.extend(fid_extension(fidapi,
batch_generate_func(
smoothed_generator,
smoothed_mapping, trainer),
seed=FLAGS.seed,
report_key='S_FID'),
trigger=(FLAGS.fid_interval, 'iteration'))
trainer.extend(
extensions.LogReport(keys=report_keys,
trigger=(FLAGS.display_interval,
'iteration')))
trainer.extend(extensions.PrintReport(report_keys),
trigger=(FLAGS.display_interval, 'iteration'))
# Recover if possible
if FLAGS.get_model_from_interation != '':
resume_iteration_str = FLAGS.get_model_from_interation
print('Resume from {}'.format(resume_iteration_str))
for model, model_name in zip(models, model_names):
chainer.serializers.load_npz(
FLAGS.out + '/' + model_name +
'_%s.npz' % resume_iteration_str,
model,
)
chainer.serializers.load_npz(
FLAGS.out + '/' + 'snapshot_iter_%s.npz' % resume_iteration_str,
trainer,
)
elif FLAGS.auto_resume:
print("Auto Resume")
candidates = []
auto_resume_dir = FLAGS.auto_resume_dir if FLAGS.auto_resume_dir != '' else FLAGS.out
for fname in [
f for f in os.listdir(auto_resume_dir)
if f.startswith('Generator_') and f.endswith('.npz')
]:
fname = re.sub(r'^Generator_', '', fname)
fname = re.sub('\.npz$', '', fname)
fname_int = None
try:
fname_int = int(fname)
except ValueError:
pass
if fname_int is not None:
all_model_exist = True
for m in model_names:
if not os.path.exists(auto_resume_dir + '/' + m + '_' +
fname + '.npz'):
all_model_exist = False
if not os.path.exists(auto_resume_dir + '/' +
('snapshot_iter_%s.npz' % fname)):
all_model_exist = False
if all_model_exist:
candidates.append(fname)
#print(candidates)
candidates.sort(key=lambda _: int(_), reverse=True)
if len(candidates) > 0:
resume_iteration_str = candidates[0]
else:
resume_iteration_str = None
if resume_iteration_str is not None:
print('Automatic resuming: use iteration %s' %
resume_iteration_str)
for model, model_name in zip(models, model_names):
chainer.serializers.load_npz(
auto_resume_dir + '/' + model_name +
'_%s.npz' % resume_iteration_str,
model,
)
chainer.serializers.load_npz(
auto_resume_dir + '/' +
'snapshot_iter_%s.npz' % resume_iteration_str,
trainer,
)
# Run the training
if FLAGS.enable_cuda_profiling:
with cupy.cuda.profile():
trainer.run()
else:
#with chainer.using_config('debug', True):
trainer.run()
for model, model_name in zip(models, model_names):
chainer.serializers.save_npz(
FLAGS.out + '/' + model_name + '_latest.npz', model)
def __init__(self):
warnings.filterwarnings('ignore')
self.start_time = time()
self.args = get_args()
if self.args.checkpoint_dir_name:
dir_name = self.args.checkpoint_dir_name
else:
dir_name = datetime.datetime.now().strftime('%y%m%d%H%M%S')
self.path_to_dir = Path(__file__).resolve().parents[1]
self.path_to_dir = os.path.join(self.path_to_dir, *['log', dir_name])
os.makedirs(self.path_to_dir, exist_ok=True)
# tensorboard
path_to_tensorboard = os.path.join(self.path_to_dir, 'tensorboard')
os.makedirs(path_to_tensorboard, exist_ok=True)
self.writer = SummaryWriter(path_to_tensorboard)
# model saving
os.makedirs(os.path.join(self.path_to_dir, 'model'), exist_ok=True)
path_to_model = os.path.join(self.path_to_dir, *['model', 'model.tar'])
# csv
os.makedirs(os.path.join(self.path_to_dir, 'csv'), exist_ok=True)
self.path_to_results_csv = os.path.join(self.path_to_dir,
*['csv', 'results.csv'])
path_to_args_csv = os.path.join(self.path_to_dir, *['csv', 'args.csv'])
if not self.args.checkpoint_dir_name:
with open(path_to_args_csv, 'a') as f:
args_dict = vars(self.args)
param_writer = csv.DictWriter(f, list(args_dict.keys()))
param_writer.writeheader()
param_writer.writerow(args_dict)
# logging by hyperdash
if not self.args.no_hyperdash:
from hyperdash import Experiment
self.exp = Experiment('Generation task on ' + self.args.dataset +
' dataset with GAN')
for key in vars(self.args).keys():
exec("self.args.%s = self.exp.param('%s', self.args.%s)" %
(key, key, key))
else:
self.exp = None
self.dataloader = get_dataloader(self.args.dataset,
self.args.image_size,
self.args.batch_size)
sample_data = self.dataloader.__iter__().__next__()[0]
image_channels = sample_data.shape[1]
z = torch.randn(self.args.batch_size, self.args.z_dim)
self.sample_z = z.view(z.size(0), z.size(1), 1, 1)
self.Generator = Generator(self.args.z_dim, image_channels,
self.args.image_size)
self.Generator_optimizer = optim.Adam(self.Generator.parameters(),
lr=self.args.lr_Generator,
betas=(self.args.beta1,
self.args.beta2))
self.writer.add_graph(self.Generator, self.sample_z)
self.Generator.to(self.args.device)
self.Discriminator = Discriminator(image_channels,
self.args.image_size)
self.Discriminator_optimizer = optim.Adam(
self.Discriminator.parameters(),
lr=self.args.lr_Discriminator,
betas=(self.args.beta1, self.args.beta2))
self.writer.add_graph(self.Discriminator, sample_data)
self.Discriminator.to(self.args.device)
self.BCELoss = nn.BCELoss()
self.sample_z = self.sample_z.to(self.args.device)
def main():
parser = argparse.ArgumentParser(
description='chainer implementation of pix2pix')
parser.add_argument('--batchsize',
'-b',
type=int,
default=1,
help='Number of images in each mini-batch')
parser.add_argument('--epoch',
'-e',
type=int,
default=200,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu',
'-g',
type=int,
default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--dataset',
'-i',
default='./facade/base',
help='Directory of image files.')
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('--seed', type=int, default=0, help='Random seed')
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')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# epoch: {}'.format(args.epoch))
print('')
# Set up a neural network to train
enc = Encoder(in_ch=12)
dec = Decoder(out_ch=3)
dis = Discriminator(in_ch=12, out_ch=3)
if args.gpu >= 0:
chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current
enc.to_gpu() # Copy the model to the GPU
dec.to_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.00001), 'hook_dec')
return optimizer
opt_enc = make_optimizer(enc)
opt_dec = make_optimizer(dec)
opt_dis = make_optimizer(dis)
train_d = FacadeDataset(args.dataset, data_range=(1, 300))
test_d = FacadeDataset(args.dataset, data_range=(300, 379))
#train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4)
#test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4)
train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize)
test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize)
# Set up a trainer
updater = FacadeUpdater(models=(enc, dec, dis),
iterator={
'main': train_iter,
'test': test_iter
},
optimizer={
'enc': opt_enc,
'dec': opt_dec,
'dis': opt_dis
},
device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
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(
enc, 'enc_iter_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dec, 'dec_iter_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dis, 'dis_iter_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.LogReport(trigger=display_interval))
trainer.extend(extensions.PrintReport([
'epoch',
'iteration',
'enc/loss',
'dec/loss',
'dis/loss',
]),
trigger=display_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(out_image(updater, enc, dec, 5, 5, args.seed, args.out),
trigger=snapshot_interval)
if args.resume:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
class Experiment():
def __init__(self, args):
self.args = args
self.writer = SummaryWriter(args.output_dir)
self.iter_i = 1
# data
transform_list = [
transforms.Resize(args.imsize),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]
train_dataset = datasets.CIFAR10(
'./data',
train=True,
transform=transforms.Compose(transform_list),
download=True)
self.train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
num_workers=args.n_workers,
shuffle=True)
# network
self.G = Generator(args.nz, args.ngf, args.nc).to(args.device)
self.D = Discriminator(args.nc, args.ndf).to(args.device)
self.criterion = nn.BCELoss()
self.optimizer_G = optim.Adam(self.G.parameters(),
lr=args.lr,
betas=(args.beta1, args.beta2))
self.optimizer_D = optim.Adam(self.D.parameters(),
lr=args.lr,
betas=(args.beta1, args.beta2))
self.real_label = 1
self.fake_label = 0
self.fixed_z = torch.randn((args.batch_size, args.nz, 1, 1),
device=args.device)
def train(self, epoch):
self.G.train()
self.D.train()
train_loss_G, train_loss_D = 0., 0.
n_samples = 0
for data, _ in self.train_loader:
batch_size = len(data)
n_samples += batch_size
real = data.to(self.args.device)
# train D
self.optimizer_D.zero_grad()
# with real
label = torch.full((batch_size, ),
self.real_label,
device=self.args.device)
output_real = self.D(real)
loss_D_real = self.criterion(output_real, label)
loss_D_real.backward()
# with fake
z = torch.randn((batch_size, self.args.nz, 1, 1),
device=self.args.device)
fake = self.G(z)
label = label.fill_(self.fake_label)
output_fake = self.D(fake)
loss_D_fake = self.criterion(output_fake, label)
loss_D_fake.backward(retain_graph=True)
loss_D = loss_D_real + loss_D_fake
self.optimizer_D.step()
# train G
self.optimizer_G.zero_grad()
label = label.fill_(self.real_label)
output_fake = self.D(fake)
loss_G = self.criterion(output_fake, label)
loss_G.backward()
self.optimizer_G.step()
loss_D = loss_D.item()
loss_G = loss_G.item()
train_loss_D += loss_D
train_loss_G += loss_G
if self.iter_i % self.args.log_freq == 0:
self.writer.add_scalar('Loss/D', loss_D, self.iter_i)
self.writer.add_scalar('Loss/G', loss_G, self.iter_i)
print('Epoch {} Train [{}/{}]: Loss/D {:.4f} Loss/G {:.4f}'.
format(epoch, n_samples, len(self.train_loader.dataset),
loss_D / batch_size, loss_G / batch_size))
self.iter_i += 1
dataset_size = len(self.train_loader.dataset)
print('Epoch {} Train: Loss/D {:.4f} Loss/G {:.4f}'.format(
epoch, train_loss_D / dataset_size, train_loss_G / dataset_size))
def test(self, epoch):
self.G.eval()
with torch.no_grad():
fake = self.G(self.fixed_z)
grid = make_grid(fake, normalize=True).cpu()
self.writer.add_image('Fake', grid, self.iter_i)
# show(grid)
fname = osp.join(self.args.output_dir,
'fake_epoch_{}.png'.format(epoch))
save_image(fake, fname, nrow=8)
def save(self, epoch):
# TODO
torch.save(self.model.state_dict(),
'./results/checkpoint_{}.pt'.format(epoch))
def run(self):
for epoch_i in range(1, 1 + self.args.epochs):
self.train(epoch_i)
self.test(epoch_i)
Args:
labels: (LongTensor) class labels, sized [N,].
num_classes: (int) number of classes.
Returns:
(tensor) encoded labels, sized [N, #classes].
"""
y = torch.eye(num_classes)
return y[labels]
# create the objects for loss function, two networks and for the two optimizers
adversarial_loss = torch.nn.BCELoss()
generator = Generator(latent=opt.latent, n_classes=opt.n_classes, num_filters=opt.num_filters, channels=opt.channels)
discriminator = Discriminator(channels=opt.channels, num_filters=opt.num_filters, n_classes=opt.n_classes)
optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.learning_rate, betas=(opt.beta_1, opt.beta_2))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.learning_rate, betas=(opt.beta_1, opt.beta_2))
# put the nets on gpu
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
generator, discriminator = generator.to(device), discriminator.to(device)
print(generator)
# generator.apply(weights_init)
# discriminator.apply(weights_init)
# start training
current_epoch = 0
for epoch in range(opt.n_epochs):
for i, (inputs, labels) in enumerate(dataloader):
inputs = inputs.to(device)
def main():
parser = argparse.ArgumentParser(description='Train GAN')
parser.add_argument('--batchsize',
'-b',
type=int,
default=64,
help='Number of images in each mini-batch')
parser.add_argument('--epoch',
'-e',
type=int,
default=100,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu',
'-g',
type=int,
default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--dataset',
'-i',
default='data/celebA/',
help='Directory of image files.')
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('--seed',
type=int,
default=0,
help='Random seed of z at visualization stage')
parser.add_argument('--snapshot_interval',
type=int,
default=10000,
help='Interval of snapshot')
parser.add_argument('--display_interval',
type=int,
default=1000,
help='Interval of displaying log to console')
parser.add_argument('--unrolling_steps', type=int, default=0)
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# batchsize: {}'.format(args.batchsize))
print('# epoch: {}'.format(args.epoch))
print('')
# Set up a neural network to train
gen = Generator()
dis = Discriminator(unrolling_steps=args.unrolling_steps)
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
gen.to_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)
# Setup a dataset
all_files = os.listdir(args.dataset)
image_files = [f for f in all_files if ('png' in f or 'jpg' in f)]
print('{} contains {} image files'.format(args.dataset, len(image_files)))
train = CelebADataset(paths=image_files, root=args.dataset)
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# Set up a trainer
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)
snapshot_interval = (args.snapshot_interval, 'iteration')
display_interval = (args.display_interval, 'iteration')
trainer.extend(extensions.snapshot(
filename='snapshot_gan_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.LogReport(trigger=display_interval,
log_name='train_gan.log'))
trainer.extend(extensions.PrintReport([
'epoch',
'iteration',
'gen/loss',
'dis/loss',
]),
trigger=display_interval)
trainer.extend(
extensions.PlotReport(['gen/loss', 'dis/loss'],
trigger=display_interval,
file_name='gan-loss.png'))
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(out_generated_image(gen, dis, 10, 10, args.seed, args.out),
trigger=snapshot_interval)
if args.resume:
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
def main():
parser = argparse.ArgumentParser(description='DCGAN')
parser.add_argument('--batchsize',
'-b',
type=int,
default=64,
help='Number of images in each mini-batch')
parser.add_argument('--epoch',
'-e',
type=int,
default=500,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu',
'-g',
type=int,
default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--dataset',
'-i',
default='',
help='Directory of image files.')
parser.add_argument('--out',
'-o',
default='result',
help='Directory to output the result image')
parser.add_argument('--resume',
'-r',
default='',
help='Resume the training from snapshot')
parser.add_argument('--gennum',
'-v',
default=10,
help='visualize image rows and columns number')
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')
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(n_hidden=args.n_hidden)
dis = Discriminator()
if args.gpu >= 0:
#modelをGPU用に変換
chainer.cuda.get_device(args.gpu).use()
gen.to_gpu()
dis.to_gpu()
#oputimizerのsetup
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)
if args.dataset == '':
#データセットの読み込み defaultはcifar10
train, _ = chainer.datasets.get_cifar10(withlabel=False, scale=255.)
else:
all_files = os.listdir(args.dataset)
image_files = [f for f in all_files if ('png' in f or 'jpg' in f)]
print('{} contains {} image files'.format(args.dataset,
len(image_files)))
train = chainer.datasets\
.ImageDataset(paths=image_files, root=args.dataset)
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
#trainerのセットアップ
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)
snapshot_interval = (args.snapshot_interval, 'iteration')
display_interval = (args.display_interval, 'iteration')
trainer.extend(
extensions.snapshot(filename='snapshot_iter_{.updater.epoch}.npz'),
trigger=(100, 'epoch'))
trainer.extend(extensions.snapshot_object(gen,
'gen_iter_{.updater.epoch}.npz'),
trigger=(100, 'epoch'))
trainer.extend(extensions.snapshot_object(dis,
'dis_iter_{.updater.epoch}.npz'),
trigger=(100, 'epoch'))
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, int(args.gennum),
int(args.gennum), args.seed, args.out),
trigger=snapshot_interval)
if args.resume:
#学習済みmodelの読み込み
chainer.serializers.load_npz(args.resume, trainer)
#学習の実行
trainer.run()
chainer.serializers.save_npz("genmodel{0}".format(args.datasets), trainer)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--batchsize', '-b', type=int, default=50)
parser.add_argument('--epoch', '-e', type=int, default=1000)
parser.add_argument('--gpu', '-g', type=int, default=-1)
parser.add_argument('--out', '-o', default='')
parser.add_argument('--resume', '-r', default='')
parser.add_argument('--n_noise', '-n', type=int, default=100)
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--snapshot_interval', type=int, default=1000)
parser.add_argument('--display_interval', type=int, default=100)
args = parser.parse_args()
out_dir = 'result'
if args.out != '':
out_dir = '{}/{}'.format(out_dir, args.out)
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# n_hidden: {}'.format(args.n_noise))
print('# epoch: {}'.format(args.epoch))
print('# out: {}'.format(out_dir))
print('')
gen = Generator(n_noise=args.n_noise, n_class=10)
dis = Discriminator(n_class=10)
if args.gpu >= 0:
cuda.get_device_from_id(args.gpu)
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
gen_optimizer = make_optimizer(gen)
dis_optimizer = make_optimizer(dis)
train, _ = chainer.datasets.get_cifar10(withlabel=True)
transformer = lambda data: (gen.make_noise(), ) + data
train = chainer.datasets.TransformDataset(train, transformer)
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
updater = CGANUpdater(models=(gen, dis),
iterator=train_iter,
optimizer={
'gen': gen_optimizer,
'dis': dis_optimizer
},
device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=out_dir)
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.LogReport(trigger=display_interval))
trainer.extend(extensions.PrintReport([
'epoch',
'iteration',
'gen/loss',
'dis/loss',
]),
trigger=display_interval)
trainer.extend(
extensions.PlotReport(['gen/loss', 'dis/loss'],
x_key='iteration',
trigger=display_interval))
trainer.extend(extensions.ProgressBar(update_interval=10))
gen_func = lambda data: gen(data[0], data[1])
def data_func(gen):
def _data_func(index):
return (gen.make_noise(), index // 10)
return _data_func
trainer.extend(
GenerateImage(gen_func,
data_func(gen),
file_name='{}/{}'.format(
out_dir, 'preview/{.updater.iteration:0>8}.png'),
rows=10,
cols=10,
seed=800,
device=args.gpu,
trigger=snapshot_interval))
trainer.run()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--workers',
type=int,
help='number of data loading workers',
default=2)
parser.add_argument('--batch',
type=int,
default=50,
help='input batch size')
parser.add_argument('--nz',
type=int,
default=100,
help='size of the latent z vector')
parser.add_argument('--ng_ch', type=int, default=64)
parser.add_argument('--nd_ch', type=int, default=64)
parser.add_argument('--epoch',
type=int,
default=50,
help='number of epochs to train for')
parser.add_argument('--lr',
type=float,
default=0.0002,
help='learning rate, default=0.0002')
parser.add_argument('--beta1',
type=float,
default=0.5,
help='beta1 for adam. default=0.5')
parser.add_argument('--outf',
default='./result',
help='folder to output images and model checkpoints')
opt = parser.parse_args()
print(opt)
batch_size = opt.batch
epoch_size = opt.epoch
try:
os.makedirs(opt.outf)
except OSError:
pass
random.seed(0)
torch.manual_seed(0)
dataset = dset.SVHN(root='../svhn_root',
download=True,
transform=transforms.Compose([
transforms.Resize(64),
transforms.ColorJitter(brightness=0,
contrast=0,
saturation=0,
hue=0.5),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
dataloader = torch.utils.data.DataLoader(dataset[:50000],
batch_size=batch_size,
shuffle=True,
num_workers=int(opt.workers))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('device:', device)
nz = int(opt.nz)
netG = Generator().to(device)
netG.apply(weights_init)
print(netG)
netD = Discriminator().to(device)
netD.apply(weights_init)
print(netD)
criterion = nn.MSELoss() # criterion = nn.BCELoss()
fixed_noise = torch.randn(batch_size, nz, 1, 1, device=device)
real_label = 1
fake_label = 0
# setup optimizer
optimizerD = optim.Adam(netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999),
weight_decay=1e-5)
optimizerG = optim.Adam(netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999),
weight_decay=1e-5)
for epoch in range(epoch_size):
for itr, data in enumerate(dataloader):
############################
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
###########################
# train with real
netD.zero_grad()
real_image = data[0].to(device)
sample_size = real_image.size(0)
label = torch.full((sample_size, ), real_label, device=device)
output = netD(real_image)
errD_real = criterion(output, label)
errD_real.backward()
D_x = output.mean().item()
# train with fake
noise = torch.randn(sample_size, nz, 1, 1, device=device)
fake_image = netG(noise)
label.fill_(fake_label)
output = netD(fake_image.detach())
errD_fake = criterion(output, label)
errD_fake.backward()
D_G_z1 = output.mean().item()
errD = errD_real + errD_fake
optimizerD.step()
############################
# (2) Update G network: maximize log(D(G(z)))
###########################
netG.zero_grad()
label.fill_(real_label) # fake labels are real for generator cost
output = netD(fake_image)
errG = criterion(output, label)
errG.backward()
D_G_z2 = output.mean().item()
optimizerG.step()
print(
'[{}/{}][{}/{}] Loss_D: {:.3f} Loss_G: {:.3f} D(x): {:.3f} D(G(z)): {:.3f}/{:.3f}'
.format(epoch + 1, opt.n_epoch, itr + 1, len(dataloader),
errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
if epoch == 0 and itr == 0:
vutils.save_image(real_image,
'{}/real_samples.png'.format(opt.outf),
normalize=True,
nrow=10)
fake_image = netG(fixed_noise)
vutils.save_image(fake_image.detach(),
'{}/fake_samples_epoch_{:03d}.png'.format(
opt.outf, epoch + 1),
normalize=True,
nrow=10)
# do checkpointing
if (epoch + 1) % 100 == 0:
torch.save(netG.state_dict(),
'{}/netG_epoch_{}.pth'.format(opt.outf, epoch + 1))
torch.save(netD.state_dict(),
'{}/netD_epoch_{}.pth'.format(opt.outf, epoch + 1))
class GAN(object):
def __init__(self):
warnings.filterwarnings('ignore')
self.start_time = time()
self.args = get_args()
if self.args.checkpoint_dir_name:
dir_name = self.args.checkpoint_dir_name
else:
dir_name = datetime.datetime.now().strftime('%y%m%d%H%M%S')
self.path_to_dir = Path(__file__).resolve().parents[1]
self.path_to_dir = os.path.join(self.path_to_dir, *['log', dir_name])
os.makedirs(self.path_to_dir, exist_ok=True)
# tensorboard
path_to_tensorboard = os.path.join(self.path_to_dir, 'tensorboard')
os.makedirs(path_to_tensorboard, exist_ok=True)
self.writer = SummaryWriter(path_to_tensorboard)
# model saving
os.makedirs(os.path.join(self.path_to_dir, 'model'), exist_ok=True)
path_to_model = os.path.join(self.path_to_dir, *['model', 'model.tar'])
# csv
os.makedirs(os.path.join(self.path_to_dir, 'csv'), exist_ok=True)
self.path_to_results_csv = os.path.join(self.path_to_dir,
*['csv', 'results.csv'])
path_to_args_csv = os.path.join(self.path_to_dir, *['csv', 'args.csv'])
if not self.args.checkpoint_dir_name:
with open(path_to_args_csv, 'a') as f:
args_dict = vars(self.args)
param_writer = csv.DictWriter(f, list(args_dict.keys()))
param_writer.writeheader()
param_writer.writerow(args_dict)
# logging by hyperdash
if not self.args.no_hyperdash:
from hyperdash import Experiment
self.exp = Experiment('Generation task on ' + self.args.dataset +
' dataset with GAN')
for key in vars(self.args).keys():
exec("self.args.%s = self.exp.param('%s', self.args.%s)" %
(key, key, key))
else:
self.exp = None
self.dataloader = get_dataloader(self.args.dataset,
self.args.image_size,
self.args.batch_size)
sample_data = self.dataloader.__iter__().__next__()[0]
image_channels = sample_data.shape[1]
z = torch.randn(self.args.batch_size, self.args.z_dim)
self.sample_z = z.view(z.size(0), z.size(1), 1, 1)
self.Generator = Generator(self.args.z_dim, image_channels,
self.args.image_size)
self.Generator_optimizer = optim.Adam(self.Generator.parameters(),
lr=self.args.lr_Generator,
betas=(self.args.beta1,
self.args.beta2))
self.writer.add_graph(self.Generator, self.sample_z)
self.Generator.to(self.args.device)
self.Discriminator = Discriminator(image_channels,
self.args.image_size)
self.Discriminator_optimizer = optim.Adam(
self.Discriminator.parameters(),
lr=self.args.lr_Discriminator,
betas=(self.args.beta1, self.args.beta2))
self.writer.add_graph(self.Discriminator, sample_data)
self.Discriminator.to(self.args.device)
self.BCELoss = nn.BCELoss()
self.sample_z = self.sample_z.to(self.args.device)
def train(self):
self.train_hist = {}
self.train_hist['Generator_loss'] = 0.0
self.train_hist['Discriminator_loss'] = 0.0
# real ---> y = 1
# fake ---> y = 0
self.y_real = torch.ones(self.args.batch_size, 1).to(self.args.device)
self.y_fake = torch.zeros(self.args.batch_size, 1).to(self.args.device)
self.Discriminator.train()
global_step = 0
# -----training -----
for epoch in range(1, self.args.n_epoch + 1):
self.Generator.train()
for idx, (x, _) in enumerate(self.dataloader):
if idx == self.dataloader.dataset.__len__(
) // self.args.batch_size:
break
z = torch.randn(self.args.batch_size, self.args.z_dim)
z = z.view(z.size(0), z.size(1), 1, 1)
z = z.to(self.args.device)
x = x.to(self.args.device)
# ----- update Discriminator -----
# minimize: -{ log[D(x)] + log[1-D(G(z))] }
self.Discriminator_optimizer.zero_grad()
# real
# ---> log[D(x)]
Discriminator_real, _ = self.Discriminator(x)
Discriminator_real_loss = self.BCELoss(Discriminator_real,
self.y_real)
# fake
# ---> log[1-D(G(z))]
Discriminator_fake, _ = self.Discriminator(self.Generator(z))
Discriminator_fake_loss = self.BCELoss(Discriminator_fake,
self.y_fake)
Discriminator_loss = Discriminator_real_loss + Discriminator_fake_loss
self.train_hist[
'Discriminator_loss'] = Discriminator_loss.item()
Discriminator_loss.backward()
self.Discriminator_optimizer.step()
# ----- update Generator -----
# As stated in the original paper,
# we want to train the Generator
# by minimizing log(1−D(G(z)))
# in an effort to generate better fakes.
# As mentioned, this was shown by Goodfellow
# to not provide sufficient gradients,
# especially early in the learning process.
# As a fix, we instead wish to maximize log(D(G(z))).
# ---> minimize: -log[D(G(z))]
self.Generator_optimizer.zero_grad()
Discriminator_fake, _ = self.Discriminator(self.Generator(z))
Generator_loss = self.BCELoss(Discriminator_fake, self.y_real)
self.train_hist['Generator_loss'] = Generator_loss.item()
Generator_loss.backward()
self.Generator_optimizer.step()
# ----- logging by tensorboard, csv and hyperdash
# tensorboard
self.writer.add_scalar('loss/Generator_loss',
Generator_loss.item(), global_step)
self.writer.add_scalar('loss/Discriminator_loss',
Discriminator_loss.item(), global_step)
# csv
with open(self.path_to_results_csv, 'a') as f:
result_writer = csv.DictWriter(
f, list(self.train_hist.keys()))
if epoch == 1 and idx == 0: result_writer.writeheader()
result_writer.writerow(self.train_hist)
# hyperdash
if self.exp:
self.exp.metric('Generator loss', Generator_loss.item())
self.exp.metric('Discriminator loss',
Discriminator_loss.item())
if (idx % 10) == 0:
self._plot_sample(global_step)
global_step += 1
elapsed_time = time() - self.start_time
print('\nTraining Finish, elapsed time ---> %f' % (elapsed_time))
def _plot_sample(self, global_step):
with torch.no_grad():
total_n_sample = min(self.args.n_sample, self.args.batch_size)
image_frame_dim = int(np.floor(np.sqrt(total_n_sample)))
samples = self.Generator(self.sample_z)
samples = samples.cpu().data.numpy().transpose(0, 2, 3, 1)
samples = (samples + 1) / 2
fig = plt.figure(figsize=(24, 15))
for i in range(image_frame_dim * image_frame_dim):
ax = fig.add_subplot(
image_frame_dim,
image_frame_dim * 2,
(int(i / image_frame_dim) + 1) * image_frame_dim + i + 1,
xticks=[],
yticks=[])
if samples[i].shape[2] == 3:
ax.imshow(samples[i])
else:
ax.imshow(samples[i][:, :, 0], cmap='gray')
self.writer.add_figure('sample images generated by GAN', fig,
global_step)
help='Random seed')
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')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# epoch: {}'.format(args.epoch))
print('')
# Set up a neural network to train
enc = Encoder(in_ch=5)
dec = Decoder(out_ch=3)
dis = Discriminator(in_ch=5, out_ch=3)
if args.gpu >= 0:
chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current
enc.to_gpu() # Copy the model to the GPU
dec.to_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.00001), 'hook_dec')
return optimizer
opt_enc = make_optimizer(enc)
opt_dec = make_optimizer(dec)
def main():
parser = argparse.ArgumentParser(
description='chainer implementation of pix2pix')
parser.add_argument('--batchsize',
'-b',
type=int,
default=1,
help='Number of images in each mini-batch')
parser.add_argument('--epoch',
'-e',
type=int,
default=500,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu',
'-g',
type=int,
default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--out',
'-o',
default='result_dehighlight',
help='Directory to output the result')
parser.add_argument('--resume',
'-r',
default='',
help='Resume the training from snapshot')
parser.add_argument('--seed', type=int, default=0, help='Random seed')
parser.add_argument('--snapshot_interval',
type=int,
default=10000,
help='Interval of snapshot')
parser.add_argument('--display_interval',
type=int,
default=10000,
help='Interval of displaying log to console')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# epoch: {}'.format(args.epoch))
print('')
# Set up a neural network to train
enc = Encoder(in_ch=3)
dec = Decoder(out_ch=1)
dis = Discriminator(in_ch=3, out_ch=1)
gen = Generator(in_ch=3, out_ch=1)
serializers.load_npz("depro.npz", depro)
gen.encoder = enc
gen.decoder = dec
if args.gpu >= 0:
chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current
enc.to_gpu() # Copy the model to the GPU
dec.to_gpu()
dis.to_gpu()
gen.to_gpu()
depro.disable_update()
# 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.00001), 'hook_dec')
return optimizer
opt_enc = make_optimizer(enc)
opt_dec = make_optimizer(dec)
opt_dis = make_optimizer(dis)
#will fix
train_d = NyuDataset("E:/nyu_depth_v2_labeled.mat",
startnum=0,
endnum=1000)
test_d = NyuDataset("E:/nyu_depth_v2_labeled.mat",
startnum=1000,
endnum=1449)
#train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4)
#test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4)
train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize)
test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize)
test_iter2 = chainer.iterators.SerialIterator(test_d,
args.batchsize,
repeat=False,
shuffle=False)
# Set up a trainer
updater = PicUpdater(models=(enc, dec, dis),
iterator={
'main': train_iter,
'test': test_iter
},
optimizer={
'enc': opt_enc,
'dec': opt_dec,
'dis': opt_dis
},
device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
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(
enc, 'enc_iter_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dec, 'dec_iter_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dis, 'dis_iter_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.LogReport(trigger=display_interval))
trainer.extend(extensions.PrintReport([
'epoch', 'iteration', 'enc/loss_enc', 'dec/loss_dec', 'dis/loss_dis',
"validation/main/loss"
]),
trigger=display_interval)
trainer.extend(extensions.Evaluator(test_iter2, gen, device=args.gpu))
trainer.extend(
extensions.PlotReport(['dec/loss_dec'],
x_key='iteration',
file_name='dec_loss.png',
trigger=display_interval))
trainer.extend(
extensions.PlotReport(['dis/loss_dis'],
x_key='iteration',
file_name='dis_loss.png',
trigger=display_interval))
trainer.extend(
extensions.PlotReport(["validation/main/loss"],
x_key='iteration',
file_name='gen_loss.png',
trigger=display_interval))
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(out_image(updater, depro, enc, dec, 3, 3, args.seed,
args.out),
trigger=snapshot_interval)
if args.resume:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
def main():
parser = argparse.ArgumentParser(description='Train script')
parser.add_argument('--data_path',
type=str,
default="data/datasets/cifar10/train",
help='dataset directory path')
parser.add_argument('--class_name',
'-class',
type=str,
default='all_class',
help='class name (default: all_class(str))')
parser.add_argument('--num_workers',
type=int,
default=4,
help='number of parallel data loading processes')
parser.add_argument('--batchsize', '-b', type=int, default=16)
parser.add_argument('--max_iter', '-m', type=int, default=40000)
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('--snapshot_interval',
type=int,
default=2500,
help='Interval of snapshot')
parser.add_argument('--evaluation_interval',
type=int,
default=5000,
help='Interval of evaluation')
parser.add_argument('--out_image_interval',
type=int,
default=1250,
help='Interval of evaluation')
parser.add_argument('--stage_interval',
type=int,
default=40000,
help='Interval of stage progress')
parser.add_argument('--display_interval',
type=int,
default=100,
help='Interval of displaying log to console')
parser.add_argument(
'--n_dis',
type=int,
default=1,
help='number of discriminator update per generator update')
parser.add_argument('--lam',
type=float,
default=10,
help='gradient penalty')
parser.add_argument('--gamma',
type=float,
default=750,
help='gradient penalty')
parser.add_argument('--pooling_comp',
type=float,
default=1.0,
help='compensation')
parser.add_argument('--pretrained_generator', type=str, default="")
parser.add_argument('--pretrained_discriminator', type=str, default="")
parser.add_argument('--initial_stage', type=float, default=0.0)
parser.add_argument('--generator_smoothing', type=float, default=0.999)
args = parser.parse_args()
record_setting(args.out)
report_keys = [
"stage", "loss_dis", "loss_gp", "loss_gen", "g", "inception_mean",
"inception_std", "FID"
]
max_iter = args.max_iter
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
generator = Generator()
generator_smooth = Generator()
discriminator = Discriminator(pooling_comp=args.pooling_comp)
# select GPU
if args.gpu >= 0:
generator.to_gpu()
generator_smooth.to_gpu()
discriminator.to_gpu()
print("use gpu {}".format(args.gpu))
if args.pretrained_generator != "":
chainer.serializers.load_npz(args.pretrained_generator, generator)
if args.pretrained_discriminator != "":
chainer.serializers.load_npz(args.pretrained_discriminator,
discriminator)
copy_param(generator_smooth, generator)
# Setup an optimizer
def make_optimizer(model, alpha=0.001, beta1=0.0, beta2=0.99):
optimizer = chainer.optimizers.Adam(alpha=alpha,
beta1=beta1,
beta2=beta2)
optimizer.setup(model)
# optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec')
return optimizer
opt_gen = make_optimizer(generator)
opt_dis = make_optimizer(discriminator)
if args.class_name == 'all_class':
data_path = args.data_path
one_class_flag = False
else:
data_path = os.path.join(args.data_path, args.class_name)
one_class_flag = True
train_dataset = ImageDataset(data_path, one_class_flag=one_class_flag)
train_iter = chainer.iterators.MultiprocessIterator(
train_dataset, args.batchsize, n_processes=args.num_workers)
# Set up a trainer
updater = Updater(models=(generator, discriminator, generator_smooth),
iterator={'main': train_iter},
optimizer={
'opt_gen': opt_gen,
'opt_dis': opt_dis
},
device=args.gpu,
n_dis=args.n_dis,
lam=args.lam,
gamma=args.gamma,
smoothing=args.generator_smoothing,
initial_stage=args.initial_stage,
stage_interval=args.stage_interval)
trainer = training.Trainer(updater, (max_iter, 'iteration'), out=args.out)
trainer.extend(extensions.snapshot_object(
generator, 'generator_{.updater.iteration}.npz'),
trigger=(args.snapshot_interval, 'iteration'))
trainer.extend(extensions.snapshot_object(
generator_smooth, 'generator_smooth_{.updater.iteration}.npz'),
trigger=(args.snapshot_interval, 'iteration'))
trainer.extend(extensions.snapshot_object(
discriminator, 'discriminator_{.updater.iteration}.npz'),
trigger=(args.snapshot_interval, 'iteration'))
trainer.extend(
extensions.LogReport(keys=report_keys,
trigger=(args.display_interval, 'iteration')))
trainer.extend(extensions.PrintReport(report_keys),
trigger=(args.display_interval, 'iteration'))
trainer.extend(sample_generate(generator_smooth, args.out),
trigger=(args.out_image_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(sample_generate_light(generator_smooth, args.out),
trigger=(args.evaluation_interval // 10, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(calc_inception(generator_smooth),
trigger=(args.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(calc_FID(generator_smooth),
trigger=(args.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(extensions.ProgressBar(update_interval=10))
# Run the training
trainer.run()
def main():
parser = argparse.ArgumentParser(description='chainer implementation of pix2pix')
parser.add_argument('--batchsize', '-b', type=int, default=1,
help='Number of images in each mini-batch')
parser.add_argument('--epoch', '-e', type=int, default=200,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu', '-g', type=int, default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--dataset', '-i', default='./facade/base',
help='Directory of image files.')
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('--seed', type=int, default=0,
help='Random seed')
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')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# epoch: {}'.format(args.epoch))
print('')
# Set up a neural network to train
enc = Encoder(in_ch=12)
dec = Decoder(out_ch=3)
dis = Discriminator(in_ch=12, out_ch=3)
if args.gpu >= 0:
chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current
enc.to_gpu() # Copy the model to the GPU
dec.to_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.00001), 'hook_dec')
return optimizer
opt_enc = make_optimizer(enc)
opt_dec = make_optimizer(dec)
opt_dis = make_optimizer(dis)
train_d = FacadeDataset(args.dataset, data_range=(1,300))
test_d = FacadeDataset(args.dataset, data_range=(300,379))
#train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4)
#test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4)
train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize)
test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize)
# Set up a trainer
updater = FacadeUpdater(
models=(enc, dec, dis),
iterator={
'main': train_iter,
'test': test_iter},
optimizer={
'enc': opt_enc, 'dec': opt_dec,
'dis': opt_dis},
device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
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(
enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
trainer.extend(extensions.LogReport(trigger=display_interval))
trainer.extend(extensions.PrintReport([
'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss',
]), trigger=display_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(
out_image(
updater, enc, dec,
5, 5, args.seed, args.out),
trigger=snapshot_interval)
if args.resume:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--workers',
type=int,
help='number of data loading workers',
default=2)
parser.add_argument('--batch_size',
type=int,
default=50,
help='input batch size')
parser.add_argument('--nz',
type=int,
default=100,
help='size of the latent z vector')
parser.add_argument('--nch_g', type=int, default=64)
parser.add_argument('--nch_d', type=int, default=64)
parser.add_argument('--n_epoch',
type=int,
default=200,
help='number of epochs to train for')
parser.add_argument('--lr',
type=float,
default=0.0002,
help='learning rate, default=0.0002')
parser.add_argument('--beta1',
type=float,
default=0.5,
help='beta1 for adam. default=0.5')
parser.add_argument('--outf',
default='./result_lsgan',
help='folder to output images and model checkpoints')
opt = parser.parse_args()
print(opt)
try:
os.makedirs(opt.outf)
except OSError:
pass
# 乱数のシード（種）を固定
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
# STL-10のトレーニングデータセットとテストデータセットを読み込む
trainset = dset.STL10(root='./dataset/stl10_root',
download=True,
split='train+unlabeled',
transform=transforms.Compose([
transforms.RandomResizedCrop(64,
scale=(88 / 96,
1.0),
ratio=(1., 1.)),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.05,
contrast=0.05,
saturation=0.05,
hue=0.05),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
])) # ラベルを使用しないのでラベルなしを混在した'train+unlabeled'を読み込む
testset = dset.STL10(root='./dataset/stl10_root',
download=True,
split='test',
transform=transforms.Compose([
transforms.RandomResizedCrop(64,
scale=(88 / 96, 1.0),
ratio=(1., 1.)),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.05,
contrast=0.05,
saturation=0.05,
hue=0.05),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
dataset = trainset + testset # STL-10のトレーニングデータセットとテストデータセットを合わせて訓練データとする
# 訓練データをセットしたデータローダーを作成する
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=int(opt.workers))
# 学習に使用するデバイスを得る。可能ならGPUを使用する
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('device:', device)
# 生成器G。ランダムベクトルから贋作画像を生成する
netG = Generator(nz=opt.nz, nch_g=opt.nch_g).to(device)
netG.apply(weights_init) # weights_init関数で初期化
print(netG)
# 識別器D。画像が、元画像か贋作画像かを識別する
netD = Discriminator(nch_d=opt.nch_d).to(device)
netD.apply(weights_init)
print(netD)
criterion = nn.MSELoss() # 損失関数は平均二乗誤差損失
# オプティマイザ−のセットアップ
optimizerD = optim.Adam(netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999),
weight_decay=1e-5) # 識別器D用
optimizerG = optim.Adam(netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999),
weight_decay=1e-5) # 生成器G用
fixed_noise = torch.randn(opt.batch_size, opt.nz, 1, 1,
device=device) # 確認用の固定したノイズ
# 学習のループ
for epoch in range(opt.n_epoch):
for itr, data in enumerate(dataloader):
real_image = data[0].to(device) # 元画像
sample_size = real_image.size(0) # 画像枚数
noise = torch.randn(sample_size, opt.nz, 1, 1,
device=device) # 正規分布からノイズを生成
real_target = torch.full((sample_size, ), 1.,
device=device) # 元画像に対する識別信号の目標値「1」
fake_target = torch.full((sample_size, ), 0.,
device=device) # 贋作画像に対する識別信号の目標値「0」
############################
# 識別器Dの更新
###########################
netD.zero_grad() # 勾配の初期化
output = netD(real_image) # 識別器Dで元画像に対する識別信号を出力
errD_real = criterion(output, real_target) # 元画像に対する識別信号の損失値
D_x = output.mean().item()
fake_image = netG(noise) # 生成器Gでノイズから贋作画像を生成
output = netD(fake_image.detach()) # 識別器Dで元画像に対する識別信号を出力
errD_fake = criterion(output, fake_target) # 贋作画像に対する識別信号の損失値
D_G_z1 = output.mean().item()
errD = errD_real + errD_fake # 識別器Dの全体の損失
errD.backward() # 誤差逆伝播
optimizerD.step() # Dのパラメーターを更新
############################
# 生成器Gの更新
###########################
netG.zero_grad() # 勾配の初期化
output = netD(fake_image) # 更新した識別器Dで改めて贋作画像に対する識別信号を出力
errG = criterion(
output, real_target) # 生成器Gの損失値。Dに贋作画像を元画像と誤認させたいため目標値は「1」
errG.backward() # 誤差逆伝播
D_G_z2 = output.mean().item()
optimizerG.step() # Gのパラメータを更新
print(
'[{}/{}][{}/{}] Loss_D: {:.3f} Loss_G: {:.3f} D(x): {:.3f} D(G(z)): {:.3f}/{:.3f}'
.format(epoch + 1, opt.n_epoch, itr + 1, len(dataloader),
errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
if epoch == 0 and itr == 0: # 初回に元画像を保存する
vutils.save_image(real_image,
'{}/real_samples.png'.format(opt.outf),
normalize=True,
nrow=10)
############################
# 確認用画像の生成
############################
fake_image = netG(fixed_noise) # 1エポック終了ごとに確認用の贋作画像を生成する
vutils.save_image(fake_image.detach(),
'{}/fake_samples_epoch_{:03d}.png'.format(
opt.outf, epoch + 1),
normalize=True,
nrow=10)
############################
# モデルの保存
############################
if (epoch + 1) % 50 == 0: # 50エポックごとにモデルを保存する
torch.save(netG.state_dict(),
'{}/netG_epoch_{}.pth'.format(opt.outf, epoch + 1))
torch.save(netD.state_dict(),
'{}/netD_epoch_{}.pth'.format(opt.outf, epoch + 1))
def main():
parser = argparse.ArgumentParser(description='Chainer example: DCGAN')
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=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--dataset',
'-i',
default='',
help='Directory of image files. Default is cifar-10.')
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')
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('')
# Set up a neural network to train
gen = Generator(n_hidden=args.n_hidden)
dis = Discriminator()
if args.gpu >= 0:
# Make a specified GPU current
chainer.backends.cuda.get_device_from_id(args.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_hooks.WeightDecay(0.0001),
'hook_dec')
return optimizer
opt_gen = make_optimizer(gen)
opt_dis = make_optimizer(dis)
if args.dataset == '':
# Load the CIFAR10 dataset if args.dataset is not specified
train, _ = chainer.datasets.get_cifar10(withlabel=False, scale=255.)
else:
all_files = os.listdir(args.dataset)
image_files = [f for f in all_files if ('png' in f or 'jpg' in f)]
print('{} contains {} image files'.format(args.dataset,
len(image_files)))
train = chainer.datasets\
.ImageDataset(paths=image_files, root=args.dataset)
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# Set up a trainer
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)
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.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=snapshot_interval)
if args.resume:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
def main():
parser = argparse.ArgumentParser(description='ChainerMN example: DCGAN')
parser.add_argument('--batchsize', '-b', type=int, default=50,
help='Number of images in each mini-batch')
parser.add_argument('--communicator', type=str,
default='hierarchical', help='Type of communicator')
parser.add_argument('--epoch', '-e', type=int, default=1000,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu', '-g', action='store_true',
help='Use GPU')
parser.add_argument('--dataset', '-i', default='',
help='Directory of image files. Default is cifar-10.')
parser.add_argument('--out', '-o', default='result',
help='Directory to output the result')
parser.add_argument('--gen_model', '-r', default='',
help='Use pre-trained generator for training')
parser.add_argument('--dis_model', '-d', default='',
help='Use pre-trained discriminator for training')
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')
args = parser.parse_args()
# Prepare ChainerMN communicator.
if args.gpu:
if args.communicator == 'naive':
print("Error: 'naive' communicator does not support GPU.\n")
exit(-1)
comm = chainermn.create_communicator(args.communicator)
device = comm.intra_rank
else:
if args.communicator != 'naive':
print('Warning: using naive communicator '
'because only naive supports CPU-only execution')
comm = chainermn.create_communicator('naive')
device = -1
if comm.rank == 0:
print('==========================================')
print('Num process (COMM_WORLD): {}'.format(comm.size))
if args.gpu:
print('Using GPUs')
print('Using {} communicator'.format(args.communicator))
print('Num hidden unit: {}'.format(args.n_hidden))
print('Num Minibatch-size: {}'.format(args.batchsize))
print('Num epoch: {}'.format(args.epoch))
print('==========================================')
# Set up a neural network to train
gen = Generator(n_hidden=args.n_hidden)
dis = Discriminator()
if device >= 0:
# Make a specified GPU current
chainer.cuda.get_device_from_id(device).use()
gen.to_gpu() # Copy the model to the GPU
dis.to_gpu()
# Setup an optimizer
def make_optimizer(model, comm, alpha=0.0002, beta1=0.5):
# Create a multi node optimizer from a standard Chainer optimizer.
optimizer = chainermn.create_multi_node_optimizer(
chainer.optimizers.Adam(alpha=alpha, beta1=beta1), comm)
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec')
return optimizer
opt_gen = make_optimizer(gen, comm)
opt_dis = make_optimizer(dis, comm)
# Split and distribute the dataset. Only worker 0 loads the whole dataset.
# Datasets of worker 0 are evenly split and distributed to all workers.
if comm.rank == 0:
if args.dataset == '':
# Load the CIFAR10 dataset if args.dataset is not specified
train, _ = chainer.datasets.get_cifar10(withlabel=False,
scale=255.)
else:
all_files = os.listdir(args.dataset)
image_files = [f for f in all_files if ('png' in f or 'jpg' in f)]
print('{} contains {} image files'
.format(args.dataset, len(image_files)))
train = chainer.datasets\
.ImageDataset(paths=image_files, root=args.dataset)
else:
train = None
train = chainermn.scatter_dataset(train, comm)
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# Set up a trainer
updater = DCGANUpdater(
models=(gen, dis),
iterator=train_iter,
optimizer={
'gen': opt_gen, 'dis': opt_dis},
device=device)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
# Some display and output extensions are necessary only for one worker.
# (Otherwise, there would just be repeated outputs.)
if comm.rank == 0:
snapshot_interval = (args.snapshot_interval, 'iteration')
display_interval = (args.display_interval, 'iteration')
# Save only model parameters.
# `snapshot` extension will save all the trainer module's attribute,
# including `train_iter`.
# However, `train_iter` depends on scattered dataset, which means that
# `train_iter` may be different in each process.
# Here, instead of saving whole trainer module, only the network models
# are saved.
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.LogReport(trigger=display_interval))
trainer.extend(extensions.PrintReport([
'epoch', 'iteration', 'gen/loss', 'dis/loss', 'elapsed_time',
]), trigger=display_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(
out_generated_image(
gen, dis,
10, 10, args.seed, args.out),
trigger=snapshot_interval)
# Start the training using pre-trained model, saved by snapshot_object
if args.gen_model:
chainer.serializers.load_npz(args.gen_model, gen)
if args.dis_model:
chainer.serializers.load_npz(args.dis_model, dis)
# Run the training
trainer.run()
def main():
parser = argparse.ArgumentParser(description='Chainer example: DCGAN')
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('--device', '-d', type=str, default='-1',
help='Device specifier. Either ChainerX device '
'specifier or an integer. If non-negative integer, '
'CuPy arrays with specified device id are used. If '
'negative integer, NumPy arrays are used')
parser.add_argument('--dataset', '-i', default='',
help='Directory of image files. Default is cifar-10.')
parser.add_argument('--out', '-o', default='result',
help='Directory to output the result')
parser.add_argument('--resume', '-r', type=str,
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')
group = parser.add_argument_group('deprecated arguments')
group.add_argument('--gpu', '-g', dest='device',
type=int, nargs='?', const=0,
help='GPU ID (negative value indicates CPU)')
args = parser.parse_args()
device = chainer.get_device(args.device)
device.use()
print('Device: {}'.format(device))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# n_hidden: {}'.format(args.n_hidden))
print('# epoch: {}'.format(args.epoch))
print('')
# Set up a neural network to train
gen = Generator(n_hidden=args.n_hidden)
dis = Discriminator()
gen.to_device(device) # Copy the model to the device
dis.to_device(device)
# 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_hooks.WeightDecay(0.0001), 'hook_dec')
return optimizer
opt_gen = make_optimizer(gen)
opt_dis = make_optimizer(dis)
if args.dataset == '':
# Load the CIFAR10 dataset if args.dataset is not specified
train, _ = chainer.datasets.get_cifar10(withlabel=False, scale=255.)
else:
all_files = os.listdir(args.dataset)
image_files = [f for f in all_files if ('png' in f or 'jpg' in f)]
print('{} contains {} image files'
.format(args.dataset, len(image_files)))
train = chainer.datasets\
.ImageDataset(paths=image_files, root=args.dataset)
# Setup an iterator
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# Setup an updater
updater = DCGANUpdater(
models=(gen, dis),
iterator=train_iter,
optimizer={
'gen': opt_gen, 'dis': opt_dis},
device=device)
# Setup a trainer
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
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.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=snapshot_interval)
if args.resume is not None:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
def main():
parser = argparse.ArgumentParser(description='Chainer example: DCGAN')
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=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--dataset',
'-i',
default='filelist.txt',
help='List file of image files.')
parser.add_argument('--dataset_root',
default='.',
help='Root directory to retrieve images from.')
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')
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))
# Set up a neural network to train
gen = Generator(n_hidden=args.n_hidden)
dis = Discriminator()
if args.gpu >= 0:
# Make a specified GPU current
chainer.cuda.get_device_from_id(args.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)
# Set up dataset
train = chainer.datasets.LabeledImageDataset(args.dataset,
root=args.dataset_root)
print('# data-size: {}'.format(len(train)))
print('# data-shape: {}'.format(train[0][0].shape))
print('')
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# Set up a trainer
updater = GANUpdater(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)
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.LogReport(trigger=display_interval))
trainer.extend(extensions.PrintReport([
'epoch',
'iteration',
'gen/loss',
'dis/loss',
]),
trigger=display_interval)
trainer.extend(
extensions.PlotReport(['gen/loss', 'dis/loss'],
x_key='iteration',
trigger=display_interval))
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(out_generated_image(gen, dis, 10, 10, args.seed, args.out),
trigger=snapshot_interval)
if args.resume:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
# Save generator model
if args.gpu >= 0:
gen.to_cpu()
chainer.serializers.save_npz(os.path.join(args.out, 'gen.npz'), gen)
