def main():
parser = argparse.ArgumentParser()
# Seed option
parser.add_argument('--seed', default=0, type=int)
# GPU option
parser.add_argument('--gpu_id', type=int, default=0)
# Dataset option
parser.add_argument('--dataset', type=str, default='cifar10to5')
# Genrator option
parser.add_argument('--g_path', type=str, required=True)
# Output options
parser.add_argument('--out', type=str, default='samples')
parser.add_argument('--num_samples', type=int, default=10)
parser.add_argument('--eval_batch_size', type=int, default=128)
args = parser.parse_args()
# Set up seed
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
# Set up GPU
if torch.cuda.is_available() and args.gpu_id >= 0:
device = torch.device('cuda:%d' % args.gpu_id)
else:
device = torch.device('cpu')
# Set up dataset
if args.dataset == 'cifar10':
vis_label_list = [[0], [1], [2], [3], [4], [5], [6], [7], [8], [9]]
elif args.dataset == 'cifar10to5':
vis_label_list = [[0], [0, 1], [1], [1, 2], [2], [2, 3], [3], [3, 4],
[4], [4, 0]]
elif args.dataset == 'cifar7to3':
vis_label_list = [[0], [0, 1], [1], [1, 2], [2], [2, 0], [0, 1, 2]]
# Set up generator
g_root = os.path.dirname(args.g_path)
g_params = util.load_params(os.path.join(g_root, 'netG_params.pkl'))
g_iteration = int(
os.path.splitext(os.path.basename(args.g_path))[0].split('_')[-1])
netG = resnet.Generator(**g_params)
netG.to(device)
netG.load_state_dict(
torch.load(args.g_path, map_location=lambda storage, loc: storage))
netG.eval()
# Set up output
if not os.path.exists(args.out):
os.makedirs(args.out)
# Set up visualizer
visualizer = Visualizer(netG, vis_label_list, device, args.out,
args.num_samples, args.eval_batch_size)
# Visualize
visualizer.visualize(g_iteration)
def main():
# For fast training
cudnn.benchmark = True
parser = argparse.ArgumentParser()
# GPU option
parser.add_argument('--gpu_id', type=int, default=0)
# Dataset options
parser.add_argument('--dataset', type=str, default='cifar10to5')
parser.add_argument('--dataroot', type=str, default='data')
parser.add_argument('--data_seed', type=int, default=12345)
parser.add_argument('--num_workers', type=int, default=4)
# Model options
parser.add_argument('--latent_dim', type=int, default=128)
parser.add_argument('--image_size', type=int, default=32)
parser.add_argument('--g_channels', type=int, default=128)
parser.add_argument('--g_spectral_norm', type=int, default=0)
parser.add_argument('--d_channels', type=int, default=128)
parser.add_argument('--d_dropout', type=int, default=1)
parser.add_argument('--d_spectral_norm', type=int, default=0)
parser.add_argument('--d_pooling', type=str, default='mean')
# Training options
parser.add_argument('--trainer', type=str, default='cpgan')
parser.add_argument('--gan_loss', type=str, default='wgan')
parser.add_argument('--batch_size', type=int, default=64)
parser.add_argument('--g_bs_multiple', type=int, default=2)
parser.add_argument('--g_lr', type=float, default=2e-4)
parser.add_argument('--d_lr', type=float, default=2e-4)
parser.add_argument('--beta1', type=float, default=0.)
parser.add_argument('--beta2', type=float, default=0.9)
parser.add_argument('--num_critic', type=int, default=5)
parser.add_argument('--lambda_gp', type=float, default=10.)
parser.add_argument('--lambda_ct', type=float, default=0.)
parser.add_argument('--lambda_cls_g', type=float, default=0.4)
parser.add_argument('--lambda_cls_d', type=float, default=1.)
parser.add_argument('--factor_m', type=float, default=0.)
parser.add_argument('--num_iterations', type=int, default=100000)
parser.add_argument('--num_iterations_decay', type=int, default=100000)
# Output options
parser.add_argument('--out', type=str, default='outputs')
parser.add_argument('--display_interval', type=int, default=100)
parser.add_argument('--snapshot_interval', type=int, default=5000)
parser.add_argument('--visualize_interval', type=int, default=5000)
parser.add_argument('--num_samples', type=int, default=10)
parser.add_argument('--eval_batch_size', type=int, default=128)
args = parser.parse_args()
args.g_spectral_norm = bool(args.g_spectral_norm)
args.d_dropout = bool(args.d_dropout)
args.d_spectral_norm = bool(args.d_spectral_norm)
# Set up GPU
if torch.cuda.is_available() and args.gpu_id >= 0:
device = torch.device('cuda:%d' % args.gpu_id)
else:
device = torch.device('cpu')
# Set up dataset
if args.dataset == 'cifar10':
args.num_classes = 10
Dataset = torchvision.datasets.CIFAR10
args.vis_label_list = [[0], [1], [2], [3], [4], [5], [6], [7], [8],
[9]]
elif args.dataset == 'cifar10to5':
args.num_classes = 5
Dataset = functools.partial(datasets.CIFAR10to5, seed=args.data_seed)
args.vis_label_list = [[0], [0, 1], [1], [1, 2], [2], [2, 3], [3],
[3, 4], [4], [4, 0]]
elif args.dataset == 'cifar7to3':
args.num_classes = 3
Dataset = functools.partial(datasets.CIFAR7to3, seed=args.data_seed)
args.vis_label_list = [[0], [0, 1], [1], [1, 2], [2], [2, 0],
[0, 1, 2]]
def normalize(x):
x = 2 * ((x * 255. / 256.) - .5)
x += torch.zeros_like(x).uniform_(0, 1. / 128)
return x
dataset = Dataset(root=args.dataroot,
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(normalize)
]))
iterator = util.InfDataLoader(dataset,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=True,
drop_last=True)
# Set up output
if not os.path.exists(args.out):
os.makedirs(args.out)
util.print_args(args, os.path.join(args.out, 'args.txt'))
# Set up models
g_params = {
'latent_dim': args.latent_dim,
'num_classes': args.num_classes,
'channels': args.g_channels,
'image_size': args.image_size,
'spectral_norm': args.g_spectral_norm
}
d_params = {
'num_classes': args.num_classes,
'channels': args.d_channels,
'dropout': args.d_dropout,
'spectral_norm': args.d_spectral_norm,
'pooling': args.d_pooling
}
netG = resnet.Generator(**g_params)
netD = resnet.ACGANDiscriminator(**d_params)
util.save_params(g_params, os.path.join(args.out, 'netG_params.pkl'))
util.save_params(d_params, os.path.join(args.out, 'netD_params.pkl'))
netG.to(device)
netD.to(device)
netG.apply(common.weights_init)
netD.apply(common.weights_init)
util.print_network(netG, 'G', os.path.join(args.out, 'netG_arch.txt'))
util.print_network(netD, 'D', os.path.join(args.out, 'netD_arch.txt'))
# Set up optimziers
optimizerG = optim.Adam(netG.parameters(),
lr=args.g_lr,
betas=(args.beta1, args.beta2))
optimizerD = optim.Adam(netD.parameters(),
lr=args.d_lr,
betas=(args.beta1, args.beta2))
# Set up learning rate schedulers
def lr_lambda(iteration):
if args.num_iterations_decay > 0:
lr = 1.0 - max(0,
(iteration + 1 -
(args.num_iterations - args.num_iterations_decay)
)) / float(args.num_iterations_decay)
else:
lr = 1.0
return lr
lr_schedulerG = optim.lr_scheduler.LambdaLR(optimizerG,
lr_lambda=lr_lambda)
lr_schedulerD = optim.lr_scheduler.LambdaLR(optimizerD,
lr_lambda=lr_lambda)
# Set up trainer
trainter_params = {
'iterator': iterator,
'models': (netG, netD),
'optimizers': (optimizerG, optimizerD),
'gan_loss': args.gan_loss,
'lr_schedulers': (lr_schedulerG, lr_schedulerD),
'batch_size': args.batch_size,
'g_bs_multiple': args.g_bs_multiple,
'num_critic': args.num_critic,
'factor_m': args.factor_m,
'device': device
}
if args.trainer == 'acgan':
Trainer = trainers.ACGANTrainer
trainter_params.update({
'lambdas': (args.lambda_gp, args.lambda_ct, args.lambda_cls_g,
args.lambda_cls_d)
})
elif args.trainer == 'cpgan':
Trainer = trainers.CPGANTrainer
trainter_params.update({
'lambdas': (args.lambda_gp, args.lambda_ct, args.lambda_cls_g,
args.lambda_cls_d)
})
trainer = Trainer(**trainter_params)
# Set up visualizer and logger
visualizer = Visualizer(netG, args.vis_label_list, device, args.out,
args.num_samples, args.eval_batch_size)
logger = Logger(args.out, 'loss')
# Train
while trainer.iteration < args.num_iterations:
iter_start_time = time.time()
trainer.update()
if (args.display_interval > 0
and trainer.iteration % args.display_interval == 0):
t = (time.time() - iter_start_time) / args.batch_size
logger.log(trainer.iteration, trainer.get_current_loss(), t)
if (args.snapshot_interval > 0
and trainer.iteration % args.snapshot_interval == 0):
torch.save(
netG.state_dict(),
os.path.join(args.out, 'netG_iter_%d.pth' % trainer.iteration))
torch.save(
netD.state_dict(),
os.path.join(args.out, 'netD_iter_%d.pth' % trainer.iteration))
if (args.visualize_interval > 0
and trainer.iteration % args.visualize_interval == 0):
visualizer.visualize(trainer.iteration)
def main(opt, reporter=None):
writer = SummaryWriter()
with open(writer.file_writer.get_logdir() + '/args.json', 'w') as f:
json.dump(opt, f)
if opt['experiment'] is None:
opt['experiment'] = 'samples'
os.system('mkdir {0}'.format(opt['experiment']))
opt['manualSeed'] = random.randint(1, 10000) # fix seed
print("Random Seed: ", opt['manualSeed'])
random.seed(opt['manualSeed'])
torch.manual_seed(opt['manualSeed'])
cudnn.benchmark = True
if torch.cuda.is_available() and not opt['cuda']:
print("WARNING: You have a CUDA device,"
"so you should probably run with --cuda")
if opt['dataset'] in ['imagenet', 'folder', 'lfw']:
# folder dataset
dataset = dset.ImageFolder(root=opt['dataroot'],
transform=transforms.Compose([
transforms.Scale(opt['imageSize']),
transforms.CenterCrop(opt['imageSize']),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
elif opt['dataset'] == 'lsun':
dataset = dset.LSUN(root=opt['dataroot'],
classes=['bedroom_train'],
transform=transforms.Compose([
transforms.Scale(opt['imageSize']),
transforms.CenterCrop(opt['imageSize']),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
elif opt['dataset'] == 'cifar10':
dataset = dset.CIFAR10(root=opt['dataroot'],
download=True,
transform=transforms.Compose([
transforms.Scale(opt['imageSize']),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
assert dataset
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=opt['batchSize'],
shuffle=True,
num_workers=int(opt['workers']))
ngpu = int(opt['ngpu'])
nz = int(opt['nz'])
ngf = int(opt['ngf'])
ndf = int(opt['ndf'])
nc = int(opt['nc'])
n_extra_layers = int(opt['n_extra_layers'])
# custom weights initialization called on netG and netD
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
if opt['noBN']:
netG = dcgan.DCGAN_G_nobn(opt['imageSize'], nz, nc, ngf, ngpu,
n_extra_layers)
elif opt['type'] == 'mlp':
netG = mlp.MLP_G(opt['imageSize'], nz, nc, ngf, ngpu)
elif opt['type'] == 'resnet':
netG = resnet.Generator(nz)
else:
netG = dcgan.DCGAN_G(opt['imageSize'], nz, nc, ngf, ngpu,
n_extra_layers)
netG.apply(weights_init)
print(netG)
if opt['type'] == 'mlp':
netD = mlp.MLP_D(opt['imageSize'], nz, nc, ndf, ngpu)
elif opt['type'] == 'resnet':
netD = resnet.Discriminator(nz)
else:
netD = dcgan.DCGAN_D(opt['imageSize'], nz, nc, ndf, ngpu,
n_extra_layers)
netD.apply(weights_init)
print(netD)
inc_noise = torch.utils.data.TensorDataset(
torch.randn(50000, nz, 1, 1).cuda())
inc_noise_dloader = torch.utils.data.DataLoader(
inc_noise, batch_size=opt['batchSize'])
input = torch.FloatTensor(opt['batchSize'], 3, opt['imageSize'],
opt['imageSize'])
noise = torch.FloatTensor(opt['batchSize'], nz, 1, 1)
fixed_noise = torch.FloatTensor(opt['batchSize'], nz, 1, 1).normal_(0, 1)
one = torch.FloatTensor([1])
mone = one * -1
if opt['cuda']:
netD.cuda()
netG.cuda()
input = input.cuda()
one, mone = one.cuda(), mone.cuda()
noise, fixed_noise = noise.cuda(), fixed_noise.cuda()
# setup optimizer
if opt['adam']:
optimizerD = optim.Adam(netD.parameters(),
lr=opt['lrD'],
betas=(opt['beta1'], opt['beta2']))
optimizerG = optim.Adam(netG.parameters(),
lr=opt['lrG'],
betas=(opt['beta1'], opt['beta2']))
else:
optimizerD = optim.RMSprop(netD.parameters(), lr=opt['lrD'])
optimizerG = optim.RMSprop(netG.parameters(), lr=opt['lrG'])
var_weight = 0.5
w = torch.tensor(
[var_weight * (1 - var_weight)**i for i in range(9, -1, -1)]).cuda()
gen_iterations = 0
for epoch in range(opt['niter']):
data_iter = iter(dataloader)
i = 0
while i < len(dataloader):
# l_var = opt.l_var + (gen_iterations + 1)/3000
l_var = opt['l_var']
############################
# (1) Update D network
###########################
for p in netD.parameters(): # reset requires_grad
p.requires_grad = True
# train the discriminator Diters times
# if gen_iterations < 25 or gen_iterations % 500 == 0:
if gen_iterations % 500 == 0:
Diters = 100
else:
Diters = opt['Diters']
j = 0
while j < Diters and i < len(dataloader):
j += 1
# enforce constraint
if not opt['var_constraint']:
for p in netD.parameters():
p.data.clamp_(opt['clamp_lower'], opt['clamp_upper'])
data = data_iter.next()
i += 1
# train with real
real_cpu, _ = data
netD.zero_grad()
batch_size = real_cpu.size(0)
if opt['cuda']:
real_cpu = real_cpu.cuda()
input.resize_as_(real_cpu).copy_(real_cpu)
inputv = Variable(input)
out_D_real = netD(inputv)
errD_real = out_D_real.mean(0).view(1)
if opt['var_constraint']:
vm_real = out_D_real.var(0)
# train with fake
noise.resize_(opt['batchSize'], nz, 1, 1).normal_(0, 1)
with torch.no_grad():
noisev = Variable(noise) # totally freeze netG
fake = netG(noisev).data
inputv = fake
out_D_fake = netD(inputv)
errD_fake = out_D_fake.mean(0).view(1)
if opt['var_constraint']:
vm_fake = out_D_fake.var(0)
errD = errD_real - errD_fake
loss = -((errD_real - errD_fake) - l_var * torch.exp(
torch.sqrt(torch.log(vm_real)**2 + torch.log(vm_fake)**2)))
loss.backward()
optimizerD.step()
if opt['var_constraint']:
writer.add_scalars('train/variance', {
'real': vm_real.item(),
'fake': vm_fake.item()
},
epoch * len(dataloader) + i)
############################
# (2) Update G network
###########################
for p in netD.parameters():
p.requires_grad = False # to avoid computation
netG.zero_grad()
# in case our last batch was the tail batch of the dataloader,
# make sure we feed a full batch of noise
noise.resize_(opt['batchSize'], nz, 1, 1).normal_(0, 1)
noisev = Variable(noise)
fake = netG(noisev)
errG = -netD(fake).mean(0).view(1)
errG.backward()
optimizerG.step()
gen_iterations += 1
if torch.isnan(errG):
raise ValueError("Loss is nan")
############################
# Log Data
###########################
print('[%d/%d][%d/%d][%d] Loss_D: %f Loss_G: %f Loss_D_real: %f'
' Loss_D_fake %f' %
(epoch, opt['niter'], i, len(dataloader), gen_iterations,
errD.data[0], errG.data[0], errD_real.data[0],
errD_fake.data[0]))
writer.add_scalar('train/critic', -errD.item(), gen_iterations)
if gen_iterations % (500 * 64 / opt['batchSize']) == 0:
real_cpu = real_cpu.mul(0.5).add(0.5)
vutils.save_image(real_cpu,
f'{opt["experiment"]}/real_samples.png')
with torch.no_grad():
fake = netG(Variable(fixed_noise))
fake.data = fake.data.mul(0.5).add(0.5)
vutils.save_image(
fake.data, f'{opt["experiment"]}/'
f'fake_samples_{gen_iterations:010d}.png')
writer.add_image(
'train/sample',
fake.data.mul(255).clamp(0, 255).byte().cpu().numpy(),
gen_iterations)
############################
# (3) Compute Scores
############################
if gen_iterations % (500 * 64 / opt['batchSize']) == 0:
with torch.no_grad():
netG.eval()
samples = []
for (x, ) in inc_noise_dloader:
samples.append(netG(x))
netG.train()
samples = torch.cat(samples, dim=0).cpu()
samples = (samples - samples.mean()) / samples.std()
score, _ = inception_score(samples.numpy(),
cuda=True,
resize=True,
splits=10)
writer.add_scalar('test/inception_50k', score, gen_iterations)
# fids = fid_score(
# samples.permute(0, 2, 3,
# 1).mul(128).add(128).clamp(255).numpy(),
# 'cifar10'
# )
# writer.add_scalar('test/fid_50k', fids, gen_iterations)
if reporter:
reporter(inception=score, fid=0)
# do checkpointing
torch.save(netG.state_dict(),
f'{opt["experiment"]}/netG_epoch_{epoch}.pth')
torch.save(netD.state_dict(),
f'{opt["experiment"]}/netD_epoch_{epoch}.pth')