def main(config):
# Initialize generator and discriminator
generator = MODELS[config.generator](2, 1)
discriminator = MODELS[config.discriminator](nc_in=1)
vgg = Vgg16(requires_grad=False)
if config.use_cuda:
device = config.gpu if config.use_cuda else None
torch.cuda.set_device(device)
vgg.cuda()
criterion = {
'criterion_GAN': torch.nn.MSELoss(),
'criterion_pixelwise': torch.nn.L1Loss(),
'L1LossMaskedMean': L1LossMaskedMean(),
'VGGLoss': VGGLoss(vgg),
'StyleLoss': StyleLoss(vgg)
}
helper = config.helper
print("helper:\t" + helper)
syn = HELPER[helper](generator, discriminator,
criterion, config)
weight_files = sorted(glob(join(syn.config.load_model_path, 'checkpoint_epoch_*.pth')), reverse=True)
print("loaded:" + weight_files[0])
syn.load_generator_history_weight(weight_files[0])
syn.move_to_cuda()
test(syn)
def __init__(self, hparams):
super().__init__()
print("Initializing our HandGAN model...")
# Workaround from https://github.com/PyTorchLightning/pytorch-lightning/issues/3998
# Happens when loading model from checkpoints. save_hyperparameters() not working
if isinstance(hparams, dict):
hparams = Namespace(**hparams)
#self.save_hyperparameters()
self.hparams = hparams
# Used to initialize the networks
init = mutils.Initializer(init_type=hparams.init_type,
init_gain=hparams.init_gain)
# Network architecture
# Two generators, one for each domain:
self.g_ab = init(generator(
hparams)) # - g_ab: translation from domain A to domain B
self.g_ba = init(generator(
hparams)) # - g_ba: translation from domain B to domain A
# Discriminators:
self.d_a = init(
discriminator(hparams)) # - d_a: domain A discriminator
self.d_b = init(
discriminator(hparams)) # - d_b: domain B discriminator
# For the perceptual discriminator we will need a feature extractor
if hparams.netD == 'perceptual':
self.vgg_net = Vgg16().eval()
# For validation we will need Inception network to compute FID metric
if hparams.valid_interval > 0.0:
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048]
self.inception_net = InceptionV3([block_idx]).eval()
if hparams.ganerated:
model = SilNet.load_from_checkpoint(
"./weights/silnet_pretrained.ckpt")
mutils.set_requires_grad(model, requires_grad=False)
self.silnet = model.unet.eval()
# ImagePool from where we randomly get generated images in both domains
self.fake_a_pool = mutils.ImagePool(hparams.pool_size)
self.fake_b_pool = mutils.ImagePool(hparams.pool_size)
# Criterions
self.crit_cycle = torch.nn.L1Loss()
self.crit_discr = DiscriminatorLoss('lsgan')
if hparams.lambda_idt > 0.0:
self.crit_idt = torch.nn.L1Loss()
if hparams.ganerated:
self.crit_geom = torch.nn.BCEWithLogitsLoss()
def main(config):
# Initialize generator and discriminator
helper = config.helper
print("helper:\t" + helper)
if 'Kumar' in helper or 'Dakai' in helper:
generator = MODELS[config.generator](n_channels=1, n_classes=1)
discriminator = MODELS[config.discriminator](in_channels=2)
else:
# Initialize generator and discriminator
generator = MODELS[config.generator](nc_in=2, nc_out=1)
discriminator = MODELS[config.discriminator](nc_in=2)
# criterion_GAN = torch.nn.L1Loss()
if 'Lcon' in helper:
vgg = Vgg16(requires_grad=False)
if config.use_cuda:
device = config.gpu if config.use_cuda else None
torch.cuda.set_device(device)
vgg.cuda()
criterion = {
'criterion_GAN': torch.nn.MSELoss(),
'criterion_pixelwise': torch.nn.L1Loss(),
'L1LossMaskedMean': L1LossMaskedMean(),
'VGGLoss': VGGLoss(vgg),
'StyleLoss': StyleLoss(vgg)
}
else:
criterion = {
'criterion_GAN': torch.nn.MSELoss(),
'criterion_pixelwise': torch.nn.L1Loss(),
'L1LossMaskedMean': L1LossMaskedMean(),
}
syn = HELPER[helper](generator, discriminator, criterion, config)
syn.move_to_cuda()
weight_files = sorted(glob(
join(syn.config.load_model_path, 'checkpoint_epoch_*.pth')),
reverse=True)
if len(weight_files) > 0:
syn.load_generator_history_weight(weight_files[0])
else:
exit(0)
test_dataset = SYNDataLoader(root=syn.config.data_root,
split_radio=syn.config.split_radio,
split='test',
data_type="SYN",
config=syn.config)
test_old(syn, test_dataset, 'final_test_epoch')
def __init__(self, opt):
"""Initialize the pix2pix class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseModel.__init__(self, opt)
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = [
'G_GAN', 'G_L1', 'G_content_low', 'G_content_deep', 'G_style',
'D_real', 'D_fake'
]
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
self.visual_names = ['real_A', 'fake_B', 'real_B']
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
if self.isTrain:
self.model_names = ['G', 'D']
else: # during test time, only load G
self.model_names = ['G']
# define networks (both generator and discriminator)
self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf,
opt.netG, opt.norm, not opt.no_dropout,
opt.init_type, opt.init_gain,
self.gpu_ids)
if self.isTrain: # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
self.netD = networks.define_D(opt.input_nc + opt.output_nc,
opt.ndf, opt.netD, opt.n_layers_D,
opt.norm, opt.init_type,
opt.init_gain, self.gpu_ids)
if self.isTrain:
# define loss functions
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
self.criterionL1 = torch.nn.L1Loss()
self.mse_loss = torch.nn.MSELoss()
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.vgg16 = Vgg16(requires_grad=False).to(device)
def main(config):
# Initialize generator and discriminator
helper = config.helper
if 'Kumar' in helper or 'Dakai' in helper:
generator = MODELS[config.generator](n_channels=1, n_classes=1)
discriminator = MODELS[config.discriminator](in_channels=2)
else:
# Initialize generator and discriminator
generator = MODELS[config.generator](nc_in=2, nc_out=1)
discriminator = MODELS[config.discriminator](nc_in=2)
vgg = Vgg16(requires_grad=False)
if config.use_cuda:
device = config.gpu if config.use_cuda else None
torch.cuda.set_device(device)
vgg.cuda()
criterion = {
'criterion_GAN': torch.nn.L1Loss(),
'criterion_pixelwise': torch.nn.L1Loss(),
'L1LossMaskedMean': L1LossMaskedMean(),
'VGGLoss': VGGLoss(vgg),
'StyleLoss': StyleLoss(vgg),
"L2LossMaskedMean": L2LossMaskedMean()
}
print("helper:\t" + helper)
syn = HELPER[helper](generator, discriminator, criterion, config)
# weight_files = sorted(glob(join(
# '../log/' + syn.config.data_name + '/' + syn.config.generator + '_' + syn.config.helper + '/' + syn.config.discriminator + '_' + str(
# syn.config.gpu) + '/checkpoint',
# 'checkpoint_epoch_*.pth')), reverse=True)
weight_files = sorted(glob(
join(syn.config.load_model_path, 'checkpoint_epoch_*.pth')),
reverse=True)
print("loaded:" + weight_files[0])
syn.load_generator_history_weight(weight_files[0])
syn.move_to_cuda()
vali_dataset = SYNDataLoader(root=syn.config.data_root,
split_radio=config.split_radio,
split='vali',
config=syn.config,
data_type="SYN")
test(syn, vali_dataset)
exit(0)
import torchvision
import pdb
def _save_img(imgs_np, file_path):
img_np = np.transpose((imgs_np[0] * 255).astype(np.uint8), (1, 2, 0))
img_pil = Image.fromarray(img_np)
img_pil.save(file_path)
return
# Resnet152 [4, 5, 6, 7]
# Vgg16 [2, 7, 14, 21, 28]
image_np = np.expand_dims(load_image(data_format='channels_first', fpath='./images/example.png', abs_path=True), axis=0)
image = numpy_to_variable(image_np)
_save_img(image_np, './temp_ori.png')
model = Vgg16()
internal = [i for i in range(29)]
attack = DispersionAttack_gpu(model, epsilon=16./255, step_size=1./255, steps=200)
adv = attack(image, attack_layer_idx_list=[14], internal=internal)
adv_np = variable_to_numpy(adv)
_save_img(adv_np, './temp_adv.png')
diff_np = np.abs(image_np - adv_np)
_save_img(diff_np, './temp_diff.png')
diff_amp_np = diff_np / diff_np.max()
_save_img(diff_amp_np, './temp_diff_amp_{0:.2f}.png'.format(1./diff_np.max()))
def run(train_loader, val_loader, epochs, lr, momentum, weight_decay, lr_step,
k1, k2, es_patience, log_dir):
model = Vgg16()
device = 'cpu'
if torch.cuda.is_available():
device = 'cuda'
model.to(device)
optimizer = optim.SGD(model.parameters(),
lr=lr,
momentum=momentum,
weight_decay=weight_decay)
lr_scheduler = ExponentialLR(optimizer, gamma=0.975)
# criterion = VAELoss(k1=k1, k2=k2).to(device)
def update_fn(engine, batch):
x, y = _prepare_batch(batch, device=device, non_blocking=True)
model.train()
optimizer.zero_grad()
output = model(x)
# Compute loss
loss = F.nll_loss(output, y)
loss.backward()
optimizer.step()
return {
"batchloss": loss.item(),
}
trainer = Engine(update_fn)
try:
GpuInfo().attach(trainer)
except RuntimeError:
print(
"INFO: By default, in this example it is possible to log GPU information (used memory, utilization). "
"As there is no pynvml python package installed, GPU information won't be logged. Otherwise, please "
"install it : `pip install pynvml`")
trainer.add_event_handler(Events.ITERATION_COMPLETED(every=lr_step),
lambda engine: lr_scheduler.step())
metric_names = [
'batchloss',
]
def output_transform(x, name):
return x[name]
for n in metric_names:
# We compute running average values on the output (batch loss) across all devices
RunningAverage(output_transform=partial(output_transform, name=n),
epoch_bound=False,
device=device).attach(trainer, n)
exp_name = datetime.now().strftime("%Y%m%d-%H%M%S")
log_path = log_dir + "/vgg_vae/{}".format(exp_name)
tb_logger = TensorboardLogger(log_dir=log_path)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=metric_names),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer, "lr"),
event_name=Events.ITERATION_STARTED)
ProgressBar(persist=True,
bar_format="").attach(trainer,
event_name=Events.EPOCH_STARTED,
closing_event_name=Events.COMPLETED)
ProgressBar(persist=False, bar_format="").attach(trainer,
metric_names=metric_names)
# val process definition
def loss_output_transform(output):
return output
def acc_output_transform(output):
return output
customed_loss = Loss(loss_fn=F.nll_loss,
output_transform=loss_output_transform,
device=device)
customed_accuracy = Accuracy(output_transform=acc_output_transform,
device=device)
metrics = {'Loss': customed_loss, 'Accuracy': customed_accuracy}
def val_update_fn(engine, batch):
model.eval()
with torch.no_grad():
x, y = _prepare_batch(batch, device=device, non_blocking=True)
output = model(x)
return output, y
val_evaluator = Engine(val_update_fn)
for name, metric in metrics.items():
metric.attach(val_evaluator, name)
def run_evaluation(engine):
val_evaluator.run(val_loader)
trainer.add_event_handler(Events.EPOCH_COMPLETED, run_evaluation)
trainer.add_event_handler(Events.COMPLETED, run_evaluation)
ProgressBar(persist=False, desc="Train evaluation").attach(val_evaluator)
# Log val metrics:
tb_logger.attach(val_evaluator,
log_handler=OutputHandler(tag="val",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
# trainer.add_event_handler(Events.ITERATION_COMPLETED, TerminateOnNan())
# Store the best model
def default_score_fn(engine):
score = engine.state.metrics['Accuracy']
return score
best_model_handler = ModelCheckpoint(dirname=log_path,
filename_prefix="best",
n_saved=3,
score_name="val_acc",
score_function=default_score_fn)
val_evaluator.add_event_handler(Events.COMPLETED, best_model_handler, {
'model': model,
})
# Add early stopping
es_patience = es_patience
es_handler = EarlyStopping(patience=es_patience,
score_function=default_score_fn,
trainer=trainer)
val_evaluator.add_event_handler(Events.COMPLETED, es_handler)
setup_logger(es_handler._logger)
setup_logger(logging.getLogger("ignite.engine.engine.Engine"))
def empty_cuda_cache(engine):
torch.cuda.empty_cache()
import gc
gc.collect()
trainer.add_event_handler(Events.EPOCH_COMPLETED, empty_cuda_cache)
val_evaluator.add_event_handler(Events.COMPLETED, empty_cuda_cache)
trainer.run(train_loader, max_epochs=epochs)
def main():
# load option
parser = TrainOptions()
opts = parser.parse()
# create output folder
if not os.path.exists(os.path.join(opts.output_dir, 'model', opts.name)):
os.mkdir(os.path.join(opts.output_dir, 'model', opts.name))
# data loader
print('\n--- load {} dataset from {} ---'.format(opts.reg_phase,
opts.dataroot))
dataset = SingleStageDataset(opts, opts.reg_phase)
loader = torch.utils.data.DataLoader(dataset,
batch_size=opts.batch_size,
shuffle=True,
num_workers=opts.nThreads,
drop_last=True)
loader_iterator = iter(cycle(loader))
dataset_other = SingleStageDataset(opts, opts.reg_phase)
loader_other = torch.utils.data.DataLoader(dataset_other,
batch_size=opts.batch_size,
shuffle=True,
num_workers=opts.nThreads,
drop_last=True)
loader_other_iterator = iter(cycle(loader_other))
# model to be regularized
opts.load = os.path.join(opts.output_dir, 'model', opts.load)
print('\n--- load {}-th stage ArtEditing model from {} ---'.format(
opts.reg_stage, opts.load))
model = Model(opts)
_, _ = model.load(opts.load)
model = model.bicycle[opts.reg_stage]
output_dim = opts.input_dim[opts.reg_stage + 1]
model.cuda()
model.eval()
# perceptual model
print('\n--- create perceptual model')
vgg = Vgg16(requires_grad=False)
vgg.cuda()
vgg.eval()
# regularizer
print('\n--- create the regularizer')
adain_optimize = AdaINOptimize(opts, model, vgg, output_dim)
if opts.reg_load != '':
opts.reg_load = os.path.join(opts.output_dir, 'model', opts.reg_load)
print(' load the regularizer from {}'.format(opts.reg_load))
ep0 = adain_optimize.load(opts.reg_load)
ep0 += 1
else:
ep0 = 0
adain_optimize.cuda()
print('start the training at epoch {}'.format(ep0 + 1))
# tensorboard
tf_board = SummaryWriter(
logdir=os.path.join(opts.output_dir, 'tfboard', opts.name))
# start the training
for it in range(ep0, opts.n_ep):
inp, out = next(loader_iterator)
# determine input output
inp = inp.cuda()
out = out.cuda()
# refine loop for this inp/out pair
_ = adain_optimize(inp, out, loader_other_iterator)
# display
adain_optimize.write_display(tf_board, it)
if (it + 1) % (opts.n_ep // 100) == 0:
print('Iteration {}/{}'.format(it + 1, opts.n_ep))
# write model file
if (it + 1) % (opts.n_ep // 10) == 0:
adain_optimize.save(
os.path.join(opts.output_dir, 'model', opts.name,
'{}.pth'.format(it + 1)), it)
tf_board.close()
return
def main(config):
# device = torch.device("cuda")
helper = config.helper
print("helper:\t" + helper)
if 'Kumar' in helper or 'Dakai' in helper:
generator = MODELS[config.generator](n_channels=1, n_classes=1)
discriminator = MODELS[config.discriminator](in_channels=2)
else:
# Initialize generator and discriminator
generator = MODELS[config.generator](nc_in=2, nc_out=1)
discriminator = MODELS[config.discriminator](nc_in=2)
# criterion_GAN = torch.nn.L1Loss()
vgg = Vgg16(requires_grad=False)
if config.use_cuda:
device = config.gpu if config.use_cuda else None
torch.cuda.set_device(device)
vgg.cuda()
criterion = {
'criterion_GAN': torch.nn.MSELoss(),
'criterion_pixelwise': torch.nn.L1Loss(),
'L1LossMaskedMean': L1LossMaskedMean(),
'VGGLoss': VGGLoss(vgg),
'StyleLoss': StyleLoss(vgg)
}
syn = HELPER[helper](generator, discriminator, criterion, config)
# if 'allGateDilateRicherContextRefine' == syn.config.generator:
# optimizer_G = Optimizer(name=syn.config.learning_algorithm,
# model=syn.generator.refine_net,
# lr=syn.config.learning_rate)
# else:
optimizer_G = Optimizer(name=syn.config.learning_algorithm,
model=syn.generator,
lr=syn.config.learning_rate)
optimizer_D = Optimizer(name=syn.config.learning_algorithm,
model=syn.discriminator,
lr=syn.config.learning_rate)
syn.move_to_cuda()
model_optimizer = {
'G': optimizer_G,
'D': optimizer_D,
}
# if 'allGateDilateRicherContextRefine' == syn.config.generator:
# syn.load_pretrained_coarse_weight()
epo = syn.load_lastest_weight(model_optimizer)
train_dataset = SYNDataLoader(root=syn.config.data_root,
split_radio=syn.config.split_radio,
split='train',
data_type="SYN",
config=syn.config)
vali_dataset = SYNDataLoader(root=syn.config.data_root,
split_radio=syn.config.split_radio,
split='vali',
data_type="SYN",
config=syn.config)
test_dataset = SYNDataLoader(root=syn.config.data_root,
split_radio=syn.config.split_radio,
split='test',
data_type="SYN",
config=syn.config)
train_loader = DataLoader(train_dataset,
batch_size=syn.config.train_batch_size,
shuffle=True,
pin_memory=True,
num_workers=syn.config.workers)
vali_loader = DataLoader(vali_dataset,
batch_size=syn.config.train_batch_size,
shuffle=False,
pin_memory=True,
num_workers=syn.config.workers)
# test_loader = DataLoader(test_dataset, batch_size=syn.config.test_batch_size, shuffle=False, pin_memory=True,
# num_workers=syn.config.workers)
# decay_epoch = [60, 90, 120, 150, 180, 210, 240]
decay_epoch = [60, 70, 80, 90, 120, 150]
decay_next = (np.array(decay_epoch) - epo > 0).argmax(axis=0)
decay_e = decay_epoch[decay_next]
for epoch in range(epo, syn.config.epochs):
# increase boundary loss factor
syn.boundary_loss_factor = min(
syn.config.boundary_loss_factor /
syn.config.max_boundary_loss_factor_epoch * epoch,
syn.config.boundary_loss_factor)
print("\nIncrease the boundary loss factor to %.3f" %
(syn.boundary_loss_factor))
train_critics = train(syn, train_loader, model_optimizer, epoch)
syn.write_summary(epoch, train_critics)
syn.plot_train_loss(epoch, train_critics)
syn.save_model_checkpoint(epoch, model_optimizer)
if (epoch + 1) % syn.config.validate_every_epoch == 0:
vali_critics = valid(syn, vali_loader, epoch)
syn.write_summary(epoch, vali_critics)
syn.plot_vali_loss(epoch, vali_critics)
# syn.write_summary(epoch, test_critics)
# decay after start_decay_at
# if (epoch + 1) % decay_e == 0:
# for g in model_optimizer['G'].param_groups:
# current_lr = max(g['lr'] * 0.5, syn.config.min_lrate)
# print("Decaying the learning ratio to %.8f" % (current_lr))
# g['lr'] = current_lr
# decay_next += 1
# decay_e = decay_epoch[decay_next]
# print("Next decay will be in the %d th epoch" % (decay_e))
test_old(syn, test_dataset, epoch)
syn.summary_writer.close()
def main():
params = {
'batch_size': 8,
'attack_model': 'inception_v3',
'input_dir_path': '/home/yantao/datasets/imagenet_100image/original',
'output_dir_path': 'images_adv',
'attack_num_steps': 1000,
'step_size': 2,
'learning_rate': 5e-2,
'attack_type': 'ori_trans', # ori/opt
}
if params['attack_model'] == 'vgg16':
params['IMAGE_SIZE'] = 224
attack_layer_idx = 14 # 0 ~ 28
model = Vgg16(attack_layer_idx=attack_layer_idx)
elif params['attack_model'] == 'resnet152':
params['IMAGE_SIZE'] = 224
attack_layer_idx = 8 # 0 ~ 8
model = Resnet152(attack_layer_idx=attack_layer_idx)
elif params['attack_model'] == 'inception_v3':
params['IMAGE_SIZE'] = 299
attack_layer_idx = [7] # 0 ~ 13
model = Inception_v3(attack_layer_idx=attack_layer_idx)
model_ori = torchvision.models.inception_v3(
pretrained=True).cuda().eval()
else:
raise ValueError('Invalid attack model type.')
model_name = model.get_name()
'''
params['output_dir_path'] = os.path.join('/home/yantao/datasets/imagenet_100image', 'DR_' + params['attack_type'] + '_' + model_name + '_layer_{0}'.format(attack_layer_idx) + '_steps_{0}_{1:01d}'.format(params['attack_num_steps'], params['step_size']))
if not os.path.exists(params['output_dir_path']):
os.mkdir(params['output_dir_path'])
'''
params['output_dir_path'] = 'images_adv'
if params['attack_type'] == 'ori':
adversary = DispersionAttack_gpu(model,
epsilon=16 / 255.,
step_size=params['step_size'] / 255.,
steps=params['attack_num_steps'])
elif params['attack_type'] == 'opt':
adversary = DispersionAttack_opt_gpu(
model,
epsilon=16 / 255.,
learning_rate=params['learning_rate'],
steps=params['attack_num_steps'])
elif params['attack_type'] == 'ori_trans':
adversary = transform_DR_attack(model,
epsilon=16 / 255.,
step_size=params['step_size'] / 255.,
steps=params['attack_num_steps'],
prob=1.0,
image_resize=330)
else:
raise ValueError('Invalid attack type.')
images_t, file_name_list = load_images(
dir_path=params['input_dir_path'],
size=[params['IMAGE_SIZE'], params['IMAGE_SIZE']],
order='channel_first',
zero_one_bound=True,
to_tensor=True)
idx = 0
pbar = tqdm(total=int(len(images_t) / params['batch_size']) + 1)
while (idx * params['batch_size'] <= len(images_t)):
if (idx + 1) * params['batch_size'] <= len(images_t):
temp_images_t = images_t[idx * params['batch_size']:(idx + 1) *
params['batch_size']]
temp_file_name_list = file_name_list[idx *
params['batch_size']:(idx +
1) *
params['batch_size']]
else:
temp_images_t = images_t[idx * params['batch_size']:]
temp_file_name_list = file_name_list[idx * params['batch_size']:]
advs_var = adversary(temp_images_t.cuda())
advs_t = advs_var.cpu()
save_images(advs_t,
dir_path=params['output_dir_path'],
file_name_list=temp_file_name_list)
pbar.update()
idx += 1
pbar.close()
def main(args=None):
if args is None:
args = sys.argv[1:]
args = parse_args(args)
args_dic = vars(args)
with open('utils/labels.txt', 'r') as inf:
args_dic['imagenet_dict'] = eval(inf.read())
args_dic['input_dir'] = os.path.join(args.dataset_dir, 'ori')
target_model = None
internal = None
attack = None
attack_layer_idx = None
if args.adv_method == 'dr':
loss_mtd = args.loss_method
if args.target_model == 'vgg16':
assert args.vgg16_attacklayer != -1
target_model = Vgg16()
internal = [i for i in range(29)]
attack_layer_idx = [args.vgg16_attacklayer] # 12, 14
args_dic['image_size'] = (224, 224)
elif args.target_model == 'resnet152':
assert args.res152_attacklayer != -1
target_model = Resnet152()
internal = [i for i in range(9)]
attack_layer_idx = [args.res152_attacklayer] # #[4, 5, 6, 7]
args_dic['image_size'] = (224, 224)
elif args.target_model == 'inception_v3':
assert args.inc3_attacklayer != -1
target_model = Inception_v3()
internal = [i for i in range(14)]
attack_layer_idx = [args.inc3_attacklayer] # [3, 4, 7, 8, 12]
args_dic['image_size'] = (299, 299)
else:
raise
attack = DispersionAttack_gpu(target_model,
epsilon=args.epsilon / 255.,
step_size=args.step_size / 255.,
steps=args.steps,
loss_mtd=loss_mtd)
elif args.adv_method == 'tidim' or args.adv_method == 'dim' or args.adv_method == 'mifgsm' or args.adv_method == 'pgd':
attack_layer_idx = [0]
internal = [0]
loss_mtd = ''
if args.target_model == 'vgg16':
target_model = torchvision.models.vgg16(
pretrained=True).cuda().eval()
args_dic['image_size'] = (224, 224)
elif args.target_model == 'resnet152':
target_model = torchvision.models.resnet152(
pretrained=True).cuda().eval()
args_dic['image_size'] = (224, 224)
elif args.target_model == 'inception_v3':
target_model = torchvision.models.inception_v3(
pretrained=True).cuda().eval()
args_dic['image_size'] = (299, 299)
else:
raise ValueError('Invalid adv_method.')
if args.adv_method == 'dim':
attack = DIM_Attack(target_model,
decay_factor=1,
prob=0.5,
epsilon=args.epsilon / 255.,
step_size=args.step_size / 255.,
steps=args.steps,
image_resize=330)
elif args.adv_method == 'mifgsm':
attack = MomentumIteratorAttack(target_model,
decay_factor=0.5,
epsilon=args.epsilon / 255.,
step_size=args.step_size / 255.,
steps=args.steps,
random_start=False)
elif args.adv_method == 'pgd':
attack = LinfPGDAttack(target_model,
epsilon=args.epsilon / 255.,
a=args.step_size / 255.,
k=args.steps,
random_start=False)
elif args.adv_method == 'tidim':
attack = TIDIM_Attack(target_model,
decay_factor=1,
prob=0.5,
epsilon=args.epsilon / 255.,
step_size=args.step_size / 255.,
steps=args.steps,
image_resize=330)
else:
raise ValueError('Invalid adv_mdthod.')
assert target_model != None and internal != None and attack != None and attack_layer_idx != None
attack_layer_idx_str = ''
for layer_idx in attack_layer_idx:
attack_layer_idx_str += (str(layer_idx) + '_')
attack_layer_idx_str = attack_layer_idx_str[:-1]
if not DEBUG:
args_dic['output_dir'] = os.path.join(
args.dataset_dir,
'{0}_{1}_layerAt_{2}_eps_{3}_stepsize_{4}_steps_{5}_lossmtd_{6}'.
format(args.adv_method, args.target_model, attack_layer_idx_str,
args.epsilon, args.step_size, args.steps, loss_mtd))
if os.path.exists(args.output_dir):
raise ValueError('Output folder existed.')
os.mkdir(args.output_dir)
count = 0
images_list = []
names_list = []
total_images = len(os.listdir(args.input_dir))
assert args.batch_size > 0
for image_count, image_name in enumerate(tqdm(os.listdir(args.input_dir))):
image_path = os.path.join(args.input_dir, image_name)
image_np = load_image(shape=args.image_size,
data_format='channels_first',
abs_path=True,
fpath=image_path)
images_list.append(image_np)
names_list.append(image_name)
count += 1
if count < args.batch_size and image_count != total_images - 1:
continue
images_np = np.array(images_list)
count = 0
images_list = []
images_var = numpy_to_variable(images_np)
if args.adv_method == 'dr':
advs = attack(images_var, attack_layer_idx, internal)
else:
assert args.batch_size == 1, 'Baselines are not tested for batch input.'
target_model.eval()
logits_nat = target_model(images_var)
y_var = logits_nat.argmax().long().unsqueeze(0)
advs = attack(images_var.cpu(), y_var.cpu())
if not DEBUG:
advs_np = variable_to_numpy(advs)
for idx, adv_np in enumerate(advs_np):
image_pil = Image.fromarray(
np.transpose((adv_np * 255).astype(np.uint8), (1, 2, 0)))
image_pil.save(
os.path.join(args.output_dir,
os.path.splitext(names_list[idx])[0] +
'.png'))
names_list = []
def __init__(self, args):
super(SDCNet2D,self).__init__()
self.rgb_max = 255
self.sequence_length = args.sequence_length
factor = 2
input_channels = self.sequence_length * 3 + (self.sequence_length - 1) * 2 #14
self.conv1 = conv2d(input_channels, 64 // factor, kernel_size=7, stride=2)
self.conv2 = conv2d(64 // factor, 128 // factor, kernel_size=5, stride=2)
self.conv3 = conv2d(128 // factor, 256 // factor, kernel_size=5, stride=2)
self.conv3_1 = conv2d(256 // factor, 256 // factor)
self.conv4 = conv2d(256 // factor, 512 // factor, stride=2)
self.conv4_1 = conv2d(512 // factor, 512 // factor)
self.conv5 = conv2d(512 // factor, 512 // factor, stride=2)
self.conv5_1 = conv2d(512 // factor, 512 // factor)
self.conv6 = conv2d(512 // factor, 1024 // factor, stride=2)
self.conv6_1 = conv2d(1024 // factor, 1024 // factor)
self.deconv5 = deconv2d(1024 // factor, 512 // factor)
self.deconv4 = deconv2d(1024 // factor, 256 // factor)
self.deconv3 = deconv2d(768 // factor, 128 // factor)
self.deconv2 = deconv2d(384 // factor, 64 // factor)
self.deconv1 = deconv2d(192 // factor, 32 // factor)
self.deconv0 = deconv2d(96 // factor, 16 // factor)
self.final_flow = nn.Conv2d(input_channels + 16 // factor, 2, kernel_size=3, stride=1, padding=1, bias=True)
# init parameters, when doing convtranspose3d, do bilinear init
for m in self.modules():
if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose3d):
if m.bias is not None:
init.uniform_(m.bias)
init.xavier_uniform_(m.weight)
self.vgg = Vgg16().type(torch.cuda.FloatTensor)
self.flownet2 = FlowNet2(args, batchNorm=False)
assert os.path.exists(args.flownet2_checkpoint), "flownet2 checkpoint must be provided."
flownet2_checkpoint = torch.load(args.flownet2_checkpoint)
self.flownet2.load_state_dict(flownet2_checkpoint['state_dict'], strict=False)
for param in self.flownet2.parameters():
param.requires_grad = False
for param in self.vgg.parameters():
param.requires_grad = False
self.warp_nn = Resample2d(bilinear=False)
self.warp_bilinear = Resample2d(bilinear=True)
self.L1Loss = nn.L1Loss()
flow_mean = torch.FloatTensor([-0.94427323, -1.23077035]).view(1, 2, 1, 1)
flow_std = torch.FloatTensor([13.77204132, 7.47463894]).view(1, 2, 1, 1)
rgb_mean = torch.FloatTensor([106.7747911, 96.13649598, 76.61428884]).view(1, 3, 1, 1)
self.register_buffer('flow_mean', flow_mean)
self.register_buffer('flow_std', flow_std)
self.register_buffer('rgb_mean', rgb_mean)
self.ignore_keys = ['flownet2', 'vgg']
return
def main():
# load option
parser = TestOptions()
opts = parser.parse()
if len(opts.gpu_ids) > 1:
raise Exception('only one GPU for testing!')
# create result folder
result_dir = os.path.join(opts.result_dir, opts.name)
if not os.path.exists(result_dir):
os.mkdir(result_dir)
# data loader
print('\n--- load {} dataset ---'.format(opts.phase))
dataset = AlignedDataset(opts)
loader = torch.utils.data.DataLoader(dataset, batch_size=opts.batch_size, shuffle=False, num_workers=opts.nThreads)
# load model
opts.load = os.path.join(opts.output_dir, 'model', opts.load)
print('\n--- load ArtEditing model from {} ---'.format(opts.load))
model = Model(opts)
model.cuda()
_, _ = model.load(opts.load)
# perceptual model
vgg = Vgg16(requires_grad=False)
vgg.cuda()
vgg.eval()
# load regularizer
print('\n--- AdaIN optimizer ---')
adain_optimize = []
for i in range(opts.n_stage - 1):
opts.reg_load[i] = os.path.join(opts.output_dir, 'model', opts.reg_load[i])
print(' load {}-th stage from {}'.format(opts.reg_load[i]))
adain_optimize.append(AdaINOptimize(opts, model.bicycle[i], vgg, opts.input_dim[i + 1]))
adain_optimize[i].cuda()
adain_optimize[i].load(opts.reg_load[i])
# test
print('\n--- Testing ---')
for idx, (imgs, imgname) in enumerate(loader):
outs = []
with torch.no_grad():
imgs = [img.cuda() for img in imgs]
# store ground-truth img at each stage
outs += imgs
names = ['gt_{}'.format(i) for i in range(opts.n_stage)]
# workflow inference
cs = model.test_forward_backward(imgs[opts.n_stage - 1])
outs += cs
names += ['infer_{}'.format(i) for i in range(len(cs))]
# artwork generation + adain optimization for each stage
zs = []
for i in range(opts.n_stage - 1):
# get input and reference output image
inp = cs[i] if i == 0 else outs[-1]
out = imgs[i + 1] if i == opts.n_stage - 2 else cs[i + 1]
# adain optimize
z, imgs_rec = adain_optimize[i].forward(inp, out)
zs.append(z)
# store reconstructed images
outs += imgs_rec
names += ['rec_{}_{:04d}'.format(i + 1, step*(opts.n_refine // 5)) for step in range(len(imgs_rec))]
# get editing results
with torch.no_grad():
# re-sample latent representations at each stage
for idx_stage in range(opts.n_stage - 1):
for idx_vary in range(opts.num):
imgs_edit = model.test_forward(cs[0], zs, vary_stage=idx_stage)
for i, img_edit in enumerate(imgs_edit):
outs.append(img_edit)
names.append('edit_{}_{}'.format(idx_stage + 1 + i, idx_stage + 1, idx_vary))
# save
save_imgs(outs, names, os.path.join(result_dir, imgname[0]))
print('{}/{}'.format(idx + 1, len(loader)))
return
def run(train_loader, val_loader, epochs, lr, momentum, log_interval, log_dir):
model = Vgg16()
writer = create_summary_writer(model, train_loader, log_dir)
device = 'cpu'
if torch.cuda.is_available():
device = 'cuda'
optimizer = optim.SGD(model.parameters(),
lr=lr,
momentum=momentum,
weight_decay=0.001)
lr_scheduler = ExponentialLR(optimizer, gamma=0.975)
trainer = create_supervised_trainer(model,
optimizer,
F.nll_loss,
device=device)
evaluator = create_supervised_evaluator(model,
metrics={
'accuracy': Accuracy(),
'nll': Loss(F.nll_loss)
},
device=device)
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda engine: lr_scheduler.step())
trainer.add_event_handler(Events.ITERATION_COMPLETED, TerminateOnNan())
# store the best model
best_model_handler = ModelCheckpoint(dirname=log_dir,
filename_prefix="best",
n_saved=3,
score_name="test_acc",
score_function=default_score_fn)
evaluator.add_event_handler(Events.COMPLETED, best_model_handler, {
'model': model,
})
# add early stopping
es_patience = 5
es_handler = EarlyStopping(patience=es_patience,
score_function=default_score_fn,
trainer=trainer)
evaluator.add_event_handler(Events.COMPLETED, es_handler)
def empty_cuda_cache(engine):
torch.cuda.empty_cache()
import gc
gc.collect()
trainer.add_event_handler(Events.EPOCH_COMPLETED, empty_cuda_cache)
evaluator.add_event_handler(Events.COMPLETED, empty_cuda_cache)
@trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
def log_training_loss(engine):
print("Epoch[{}] Iteration[{}/{}] Loss: {:.2f}"
"".format(engine.state.epoch, engine.state.iteration,
len(train_loader), engine.state.output))
writer.add_scalar("training/loss", engine.state.output,
engine.state.iteration)
@trainer.on(Events.EPOCH_COMPLETED)
def log_training_results(engine):
evaluator.run(train_loader)
metrics = evaluator.state.metrics
avg_accuracy = metrics['accuracy']
avg_nll = metrics['nll']
print(
"Training Results - Epoch: {} Avg accuracy: {:.2f} Avg loss: {:.2f}"
.format(engine.state.epoch, avg_accuracy, avg_nll))
writer.add_scalar("training/avg_loss", avg_nll, engine.state.epoch)
writer.add_scalar("training/avg_accuracy", avg_accuracy,
engine.state.epoch)
@trainer.on(Events.EPOCH_COMPLETED)
def log_validation_results(engine):
evaluator.run(val_loader)
metrics = evaluator.state.metrics
avg_accuracy = metrics['accuracy']
avg_nll = metrics['nll']
print(
"Validation Results - Epoch: {} Avg accuracy: {:.2f} Avg loss: {:.2f}"
.format(engine.state.epoch, avg_accuracy, avg_nll))
writer.add_scalar("valdation/avg_loss", avg_nll, engine.state.epoch)
writer.add_scalar("valdation/avg_accuracy", avg_accuracy,
engine.state.epoch)
# kick everything off
trainer.run(train_loader, max_epochs=epochs)
writer.close()
