def eval():
model.eval()
count = 1
avg_psnr_predicted = 0.0
avg_ssim_predicted = 0.0
avg_time_predicted = 0.0
# batch_size = opt.testBatchSize
for batch in testing_data_loader:
input, targets, neigbor, flow, bicubic, filename = batch[0], batch[
1], batch[2], batch[3], batch[4], batch[5]
with torch.no_grad():
input = Variable(input).cuda(gpus_list[0])
bicubic = Variable(bicubic).cuda(gpus_list[0])
neigbor = [Variable(j).cuda(gpus_list[0]) for j in neigbor]
flow = [Variable(j).cuda(gpus_list[0]).float() for j in flow]
t0 = time.time()
if opt.chop_forward:
with torch.no_grad():
predictions = chop_forward(input, neigbor, flow, model,
opt.upscale_factor)
else:
with torch.no_grad():
predictions = model(input, neigbor, flow)
t1 = time.time()
time_predicted = (t1 - t0) / len(predictions)
for i in range(len(predictions)):
cur_video = filename[i].split('/')[-2]
cur_frame = filename[i].split('/')[-1]
prediction = utils.tensor2img(predictions[i])
save_img(prediction, cur_video + '_' + cur_frame, False)
avg_time_predicted += time_predicted
count += 1
if opt.HR_dir != '':
target = utils.tensor2img(targets[i])
psnr_predicted = PSNR(target, prediction)
ssim_predicted = SSIM(target, prediction)
avg_psnr_predicted += psnr_predicted
avg_ssim_predicted += ssim_predicted
logger.info(
"Processing: %s || PSNR: %.4f || SSIM: %.4f || Timer: %.4f sec."
% (cur_video + '_' + cur_frame, psnr_predicted,
ssim_predicted, time_predicted))
else:
logger.info("Processing: %s || Timer: %.4f sec." %
(cur_video + '_' + cur_frame, time_predicted))
if opt.HR_dir != '':
logger.info("AVG_PSNR = {}".format(avg_psnr_predicted / count))
logger.info("AVG_SSIM = {}".format(avg_ssim_predicted / count))
logger.info("AVG_TIME = {}".format(avg_time_predicted / count))
def gen_adv():
mse = 0
original_files = get_original_file(input_dir + val_list)
#下一个图片的初始梯度方向为上一代的最后的值
global momentum
momentum = 0
for filename, label in original_files:
img_path = input_dir + filename
print("Image: {0} ".format(img_path))
img = process_img(img_path)
#adv_img = attack_nontarget_by_ensemble(img, label,origdict[label],label)
adv_img, m = attack_nontarget_by_ensemble(img, label, origdict[label],
label, momentum)
#m为上一个样本最后一次梯度值
momentum = m
#adv_img 已经经过转换了，范围是0-255
image_name, image_ext = filename.split('.')
##Save adversarial image(.png)
save_adv_image(adv_img, output_dir + image_name + '.png')
org_img = tensor2img(img)
score = calc_mse(org_img, adv_img)
print("Image:{0}, mase = {1} ".format(img_path, score))
mse += score
print("ADV {} files, AVG MSE: {} ".format(len(original_files),
mse / len(original_files)))
def gen_adv():
mse = 0
adv_acc = 0
original_files = get_original_file(input_dir + val_list)
#init_files = get_init_file(init_dir+'init_list.txt')
for idx, (filename,
label) in enumerate(original_files[args.start:args.end]):
img_path = input_dir + filename
#init_path = init_dir + init_files[idx][0] + '.jpg'
image_name, image_ext = filename.split('.')
if image_name in area_rank[:40]:
SPARSE_PER = 99
if image_name in area_rank[40:80]:
SPARSE_PER = 97
if image_name in area_rank[80:]:
SPARSE_PER = 95
#bboxes = get_bbox('mask/'+image_name+'.xml')
bboxes = None
if verbose:
print("Image: {0} ".format(img_path))
img = process_img(img_path)
#init = process_img(init_path) * 0.01
init = None
adv_img, adv_label = attack_by_MPGD(img, label, bboxes, SPARSE_PER,
init)
save_adv_image(adv_img, output_dir + image_name + '.png')
org_img = tensor2img(img)
score = calc_mse(org_img, adv_img)
mse += score if label != adv_label else 128
adv_acc += 1 if label == adv_label else 0
if label == adv_label:
print("model: ", i, "\timage: ", filename, label)
print("ADV {} files, AVG MSE: {}, ADV_ACC: {} ".\
format(len(original_files), mse/len(original_files),adv_acc))
def attack_nontarget_by_PGD(img, src_label):
pred_label = src_label
step = 8.0 / 256.0
eps = 16.0 / 256.0
while pred_label == src_label:
#生成对抗样本
adv = PGD(adv_program=adv_program,
eval_program=eval_program,
gradients=gradients,
o=img,
input_layer=input_layer,
output_layer=out,
step_size=step,
epsilon=eps,
iteration=10,
isTarget=False,
target_label=0,
use_gpu=use_gpu)
pred_label, pred_score = inference(adv)
step *= 2
if step > eps:
break
print("Test-score: {0}, class {1}".format(pred_score, pred_label))
adv_img = tensor2img(adv)
return adv_img
def attack_nontarget_by_PGD(adv_prog, img, pred_label, src_label, out=None):
# pred_label = [src_label, src_label]
step = 8.0 / 256.0
eps = 128.0 / 256.0
while src_label in pred_label:
# 生成对抗样本
adv = T_PGD(adv_program=adv_prog, eval_program=double_eval_program, gradients=gradients, o=img,
input_layer=input_layer, output_layer=out, step_size=step, epsilon=eps, iteration=10,
t = 0.6, pix_num=224*224*3/30, isTarget=False, target_label=0, use_gpu=use_gpu)
pred_label, pred_score = inference(adv, out)
print("the current label is {}".format(pred_label))
print("the current step is {}".format(step))
#step += 1.0 / 256.0
step *= 1.5
if step > eps:
break
print("Test-score: {0}, class {1}".format(pred_score, pred_label))
adv_img = tensor2img(adv)
return adv_img
def attack_nontarget_by_FGSM(img, src_label, target):
pred_label = src_label
step = float(args.step_size)
eps = float(args.eps)
while pred_label == src_label:
#生成对抗样本
adv = PGDL2(adv_program=adv_program,
eval_program=eval_program,
loss=loss,
gradients=gradients,
o=img,
input_layer=input_layer,
output_layer=out_logits,
input_layer_eval=input_layer_eval,
output_layer_eval=out_eval,
step_size=step,
epsilon=eps,
iteration=int(args.eps / step),
isTarget=IsTarget,
target_label=target,
src_label=src_label,
use_gpu=True,
discrete=False,
confidence=args.confidence)
pred_label, pred_score, _ = inference(adv)
step *= 2
if step > eps:
break
print("Test-score: {0}, class {1}".format(pred_score, pred_label))
adv_img = tensor2img(adv)
return adv_img
def gen_adv():
mse = 0
original_files = get_original_file(input_dir + val_list)
num = 1
cout = 0
print("the model is {}".format(model_name))
for filename, label in original_files:
img_path = input_dir + filename
print("Image: {0} ".format(img_path))
img = process_img(img_path)
#print(img.shape)
result = exe.run(eval_program,
fetch_list=[out],
feed={input_layer.name: img})
result = result[0][0]
o_label = np.argsort(result)[::-1][:1][0]
print("原始标签为{0}".format(o_label))
if o_label == int(label):
adv_img = attack_nontarget_by_FGSM(img, label)
#adv_img = attack_nontarget_by_PGD(img, label)
else:
print("{0}个样本已为对抗样本, name为{1}".format(num, filename))
img = tensor2img(img)
#print(img.shape)
image_name, image_ext = filename.split('.')
save_adv_image(img, output_dir + image_name + '.png')
num += 1
cout += 1
continue
image_name, image_ext = filename.split('.')
##Save adversarial image(.png)
save_adv_image(adv_img, output_dir + image_name + '.png')
org_img = tensor2img(img)
score = calc_mse(org_img, adv_img)
mse += score
num += 1
print("成功attack的有 {}".format(120 - cout))
print("ADV {} files, AVG MSE: {} ".format(len(original_files),
mse / len(original_files)))
def gen_adv():
mse = 0
original_files = get_original_file(input_dir + val_list)
for filename, label in original_files:
img_path1 = input_dir + filename
img_path2 = output_dir + filename.split('.')[0] + '.png'
print("Image: {0} ".format(img_path1))
img1=process_img(img_path1)
img2=process_img(img_path2)
img1 = tensor2img(img1)
img2 = tensor2img(img2)
score = calc_mse(img1, img2)
mse += score
print("ADV {} files, AVG MSE: {} ".format(len(original_files), mse/len(original_files)))
def attack_driver(img, src_label, filename):
eps = 16. / 255.
step = eps / 50
iteration = 300
imgname, ext = os.path.splitext(filename)
adv = MIFGSM(adv_program=adv_program, o=img,
input_layer=input_layer, step_size=step, epsilon=eps,
gt_label=src_label, use_gpu=use_gpu,
iteration=iteration, gradients=gradients, imgname=imgname)
print("{1}\ttarget class {0}".format(src_label, filename))
adv_img = tensor2img(adv)
return adv_img
def attack_nontarget_by_FGSM(img, src_label):
pred_label = src_label
#mom = 0.8
step = 8.0 / 256.0
eps = 128.0 / 256.0
while pred_label == src_label:
#生成对抗样本
# adv=L_PGD(adv_program=adv_program,eval_program=eval_program,gradients=gradients,o=img,
# input_layer=input_layer,output_layer=out, step_size=step,epsilon=eps, iteration=8, pix_num=224*224*3/30,
# isTarget=False,target_label=0,use_gpu=use_gpu)
# adv = G_FGSM(adv_program=adv_program, eval_program=eval_program, gradients=gradients, o=img,
# input_layer=input_layer, output_layer=out, step_size=step, epsilon=eps,
# pix_num=224 * 224 * 3, isTarget=False, target_label=0, use_gpu=use_gpu)
# adv = T_PGD(adv_program=adv_program, eval_program=eval_program, gradients=gradients, o=img,
# input_layer=input_layer, output_layer=out, step_size=step, epsilon=eps,iteration=8,t=0,
# pix_num=3*224*224/30, isTarget=False, target_label=0, use_gpu=use_gpu)
adv = T_FGSM(adv_program=adv_program,
eval_program=eval_program,
gradients=gradients,
o=img,
input_layer=input_layer,
output_layer=out,
step_size=step,
epsilon=eps,
t=0.6,
pix_num=3 * 224 * 224 / 30,
isTarget=False,
target_label=0,
use_gpu=use_gpu)
pred_label, pred_score = inference(adv)
step *= 1.5
#mom *= 0.8
if step > eps:
break
print("Test-score: {0}, class {1}".format(pred_score, pred_label))
if pred_label != src_label:
print("攻击成功，{0}->{1}".format(src_label, pred_label))
else:
adv = img
print("攻击失败，去除扰动，保存为源图像")
adv_img = tensor2img(adv)
return adv_img
def gen_adv():
mse = 0
original_files = get_original_file(input_dir + val_list)
for filename, gt_label in original_files:
img_path = input_dir + filename
img = process_img(img_path)
image_name, image_ext = filename.split('.')
adv_img = attack_driver(img, gt_label, filename)
save_adv_image(adv_img, output_dir + image_name + '.png')
org_img = tensor2img(img)
score = calc_mse(org_img, adv_img)
print(score)
mse += score
print("ADV {} files, AVG MSE: {} ".format(len(original_files), mse / len(original_files)))
def gen_adv():
mse = 0
original_files = get_original_file('input_image/' + val_list)
for filename, label in original_files:
img_path = input_dir + filename.split('.')[0] + '.png'
print("Image: {0} ".format(img_path))
img = process_img(img_path)
adv_img = attack_nontarget_by_SINIFGSM(img, label)
image_name, image_ext = filename.split('.')
# Save adversarial image(.png)
save_adv_image(adv_img, output_dir + image_name + '.png')
org_img = tensor2img(img)
score = calc_mse(org_img, adv_img)
mse += score
print("ADV {} files, AVG MSE: {} ".format(len(original_files),
mse / len(original_files)))
def evaluate_save(model_G, model_FRCNN, data_loader, device, config):
i = 0
print("Detection started........")
for image, targets in data_loader:
image['image_lq'] = image['image_lq'].to(device)
_, img, _, _ = model_G(image['image_lq'])
img_count = img.size()[0]
images = [img[i] for i in range(img_count)]
outputs = model_FRCNN(images)
file_name = os.path.splitext(os.path.basename(image['LQ_path'][0]))[0]
file_path = os.path.join(config['path']['Test_Result_SR'],
file_name + '.txt')
i = i + 1
print(i)
img = img[0].detach()[0].float().cpu()
img = tensor2img(img)
get_prediction(outputs, file_path, config, file_name, img)
print("successfully generated the results!")
def attack_nontarget_by_SINIFGSM(img, src_label):
pred_label, pred_score, pred1, pred2, pred3, pred4, pred5, pred6, pred7, pred8 = inference(
img)
step = 8.0 / 256.0
eps = 32.0 / 256.0
adv = img
label_lists = [
pred_label, pred1, pred2, pred3, pred4, pred5, pred6, pred7, pred8
]
m = 0
y = np.array([2, 1, 1, 1, 1, 1, 1, 1])
while src_label in label_lists[:7]:
m = m + 1
for i in range(len(label_lists) - 1):
if label_lists[i + 1] == src_label:
y[i] = 3
y[7] = 1
adv = SI_NI_FGSM(o=adv,
y=y,
mlabel=pred_label,
step_size=step,
epsilon=eps,
isTarget=False,
target_label=0,
use_gpu=use_gpu,
T=5,
u=0.8)
label_lists[0], pred_score, label_lists[1], label_lists[
2], label_lists[3], label_lists[4], label_lists[5], label_lists[
6], label_lists[7], label_lists[8] = inference(adv)
print('labels:{0} {1} {2} {3} {4} {5} {6} {7} {8} {9}'.format(
src_label, label_lists[0], label_lists[1], label_lists[2],
label_lists[3], label_lists[4], label_lists[5], label_lists[6],
label_lists[7], label_lists[8]))
if m > 10:
break
print("Test-score: {0}, class {1}".format(pred_score, pred_label))
adv_img = tensor2img(adv)
return adv_img
def gen_adv():
mse = 0
original_files = get_original_file('./input_image/' + val_list)
target_label_list = [
76, 18, 104, 36, 72, 72, 47, 92, 113, 5, 84, 74, 82, 34, 42, 84, 70,
98, 29, 87, 104, 94, 103, 61, 21, 83, 108, 104, 26, 112, 84, 107, 104,
45, 72, 19, 72, 75, 55, 104, 54, 104, 72, 74, 91, 25, 68, 107, 91, 41,
116, 21, 104, 56, 102, 51, 46, 87, 113, 19, 113, 85, 24, 93, 110, 102,
24, 84, 27, 38, 48, 43, 10, 32, 68, 87, 54, 12, 84, 29, 3, 13, 26, 2,
3, 106, 105, 34, 118, 66, 19, 74, 63, 42, 9, 113, 21, 6, 40, 40, 21,
104, 86, 23, 40, 12, 37, 20, 40, 12, 79, 15, 9, 48, 74, 51, 91, 79, 46,
80
]
# hard examples need use targeted attack
for filename, label in original_files[args.start_idx:args.end_idx]:
img_path = input_dir + filename.split('.')[0] + args.subfix
print("Image: {0} ".format(img_path))
img = process_img(img_path)
target = target_label_list[label - 1]
if IsTarget:
print('target class', target)
adv_img = attack_nontarget_by_FGSM(img, label, target)
# adv_img = attack_nontarget_by_FGSM(img, label)
image_name, image_ext = filename.split('.')
##Save adversarial image(.png)
save_adv_image(adv_img, output_dir + image_name + '.png')
org_img = tensor2img(img)
score = calc_mse(org_img, adv_img)
mse += score
print('MSE %.2f' % (score))
sys.stdout.flush()
print("ADV {} files, AVG MSE: {} ".format(len(original_files),
mse / len(original_files)))
def attack_nontarget_by_ensemble(img, src_label, src_label2, label,
momentum): #src_label2为转换后的标签
adv, m = ensem_mom_attack_threshold_9model_tarversion(
adv_program=adv_program,
eval_program=eval_program,
gradients=gradients,
o=img,
src_label=src_label,
src_label2=src_label2,
label=label,
out1=out1,
out2=out2,
out3=out3,
out4=out4,
out5=out5,
out6=out6,
out7=out7,
out8=out8,
out9=out9,
mm=momentum) #添加了mm
adv_img = tensor2img(adv)
return adv_img, m
def attack_by_MPGD(src_img, src_label, bboxes, sparse_per, init=None):
verbose = args.verbose
step_size = args.step_size
epsilon = args.clip_eps
attack_iters = args.attack_iters
attack_confidence = args.attack_confidence
sparse_percentage = args.sparse_percentage
#sparse_percentage = sparse_per
momentum_decay = args.momentum_decay
lr_decay = args.lr_decay
init = init
adv = MPGD(adv_program,
gradients,
src_img,
src_label,
input_layer,
noise_layer,
out,
loss,
loss_cls,
bboxes,
step_size=step_size,
epsilon=epsilon,
iteration=attack_iters,
lr_decay=lr_decay,
use_gpu=use_gpu,
confidence=attack_confidence,
verbose=verbose,
init=init,
decay_factor=momentum_decay,
diversity_iter=diversity_iter,
sparse_percentage=sparse_percentage,
gradient_sign=args.gradient_sign,
norm_regulizer=args.norm_regulizer)
adv_img = tensor2img(adv)
pred_label, pred_score = inference(adv)
return adv_img, pred_label
### load occlusion mask
filename = os.path.join(occ_dir, "%05d.png" % (t - 1))
occ_mask = utils.read_img(filename)
noc_mask = 1 - occ_mask
with torch.no_grad():
## convert to tensor
img2 = utils.img2tensor(img2).to(device)
flow = utils.img2tensor(flow).to(device)
## warp img2
warp_img2 = flow_warping(img2, flow)
## convert to numpy array
warp_img2 = utils.tensor2img(warp_img2)
## compute warping error
diff = np.multiply(warp_img2 - img1, noc_mask)
N = np.sum(noc_mask)
if N == 0:
N = diff.shape[0] * diff.shape[1] * diff.shape[2]
err += np.sum(np.square(diff)) / N
err_all[v] = err / (len(frame_list) - 1)
print("\nAverage Warping Error = %f\n" % (err_all.mean()))
err_all = np.append(err_all, err_all.mean())
def main():
#### setup options of three networks
parser = argparse.ArgumentParser()
parser.add_argument("-opt", type=str, help="Path to option YMAL file.")
parser.add_argument(
"--launcher", choices=["none", "pytorch"], default="none", help="job launcher"
)
parser.add_argument("--local_rank", type=int, default=0)
args = parser.parse_args()
opt = option.parse(args.opt, is_train=True)
# convert to NoneDict, which returns None for missing keys
opt = option.dict_to_nonedict(opt)
# choose small opt for SFTMD test, fill path of pre-trained model_F
#### set random seed
# seed = opt["train"]["manual_seed"]
# if seed is None:
# seed = random.randint(1, 10000)
# load PCA matrix of enough kernel
print("load PCA matrix")
pca_matrix = torch.load(
opt["pca_matrix_path"], map_location=lambda storage, loc: storage
)
print("PCA matrix shape: {}".format(pca_matrix.shape))
#### distributed training settings
if args.launcher == "none": # disabled distributed training
opt["dist"] = False
opt["dist"] = False
rank = -1
print("Disabled distributed training.")
else:
opt["dist"] = True
opt["dist"] = True
init_dist()
world_size = (
torch.distributed.get_world_size()
) # Returns the number of processes in the current process group
rank = torch.distributed.get_rank() # Returns the rank of current process group
util.set_random_seed(0)
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
###### Predictor&Corrector train ######
#### loading resume state if exists
if opt["path"].get("resume_state", None):
# distributed resuming: all load into default GPU
device_id = torch.cuda.current_device()
resume_state = torch.load(
opt["path"]["resume_state"],
map_location=lambda storage, loc: storage.cuda(device_id),
)
option.check_resume(opt, resume_state["iter"]) # check resume options
else:
resume_state = None
#### mkdir and loggers
if rank <= 0: # normal training (rank -1) OR distributed training (rank 0-7)
if resume_state is None:
# Predictor path
util.mkdir_and_rename(
opt["path"]["experiments_root"]
) # rename experiment folder if exists
util.mkdirs(
(
path
for key, path in opt["path"].items()
if not key == "experiments_root"
and "pretrain_model" not in key
and "resume" not in key
)
)
os.system("rm ./log")
os.symlink(os.path.join(opt["path"]["experiments_root"], ".."), "./log")
# config loggers. Before it, the log will not work
util.setup_logger(
"base",
opt["path"]["log"],
"train_" + opt["name"],
level=logging.INFO,
screen=False,
tofile=True,
)
util.setup_logger(
"val",
opt["path"]["log"],
"val_" + opt["name"],
level=logging.INFO,
screen=False,
tofile=True,
)
logger = logging.getLogger("base")
logger.info(option.dict2str(opt))
# tensorboard logger
if opt["use_tb_logger"] and "debug" not in opt["name"]:
version = float(torch.__version__[0:3])
if version >= 1.1: # PyTorch 1.1
from torch.utils.tensorboard import SummaryWriter
else:
logger.info(
"You are using PyTorch {}. Tensorboard will use [tensorboardX]".format(
version
)
)
from tensorboardX import SummaryWriter
tb_logger = SummaryWriter(log_dir="log/{}/tb_logger/".format(opt["name"]))
else:
util.setup_logger(
"base", opt["path"]["log"], "train", level=logging.INFO, screen=False
)
logger = logging.getLogger("base")
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
#### create train and val dataloader
dataset_ratio = 200 # enlarge the size of each epoch
for phase, dataset_opt in opt["datasets"].items():
if phase == "train":
train_set = create_dataset(dataset_opt)
train_size = int(math.ceil(len(train_set) / dataset_opt["batch_size"]))
total_iters = int(opt["train"]["niter"])
total_epochs = int(math.ceil(total_iters / train_size))
if opt["dist"]:
train_sampler = DistIterSampler(
train_set, world_size, rank, dataset_ratio
)
total_epochs = int(
math.ceil(total_iters / (train_size * dataset_ratio))
)
else:
train_sampler = None
train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler)
if rank <= 0:
logger.info(
"Number of train images: {:,d}, iters: {:,d}".format(
len(train_set), train_size
)
)
logger.info(
"Total epochs needed: {:d} for iters {:,d}".format(
total_epochs, total_iters
)
)
elif phase == "val":
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt, opt, None)
if rank <= 0:
logger.info(
"Number of val images in [{:s}]: {:d}".format(
dataset_opt["name"], len(val_set)
)
)
else:
raise NotImplementedError("Phase [{:s}] is not recognized.".format(phase))
assert train_loader is not None
assert val_loader is not None
#### create model
model = create_model(opt) # load pretrained model of SFTMD
#### resume training
if resume_state:
logger.info(
"Resuming training from epoch: {}, iter: {}.".format(
resume_state["epoch"], resume_state["iter"]
)
)
start_epoch = resume_state["epoch"]
current_step = resume_state["iter"]
model.resume_training(resume_state) # handle optimizers and schedulers
else:
current_step = 0
start_epoch = 0
prepro = util.SRMDPreprocessing(
scale=opt["scale"], pca_matrix=pca_matrix, cuda=True, **opt["degradation"]
)
kernel_size = opt["degradation"]["ksize"]
padding = kernel_size // 2
#### training
logger.info(
"Start training from epoch: {:d}, iter: {:d}".format(start_epoch, current_step)
)
for epoch in range(start_epoch, total_epochs + 1):
if opt["dist"]:
train_sampler.set_epoch(epoch)
for _, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
LR_img, ker_map, kernels = prepro(train_data["GT"], True)
LR_img = (LR_img * 255).round() / 255
model.feed_data(
LR_img, GT_img=train_data["GT"], ker_map=ker_map, kernel=kernels
)
model.optimize_parameters(current_step)
model.update_learning_rate(
current_step, warmup_iter=opt["train"]["warmup_iter"]
)
visuals = model.get_current_visuals()
if current_step % opt["logger"]["print_freq"] == 0:
logs = model.get_current_log()
message = "<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> ".format(
epoch, current_step, model.get_current_learning_rate()
)
for k, v in logs.items():
message += "{:s}: {:.4e} ".format(k, v)
# tensorboard logger
if opt["use_tb_logger"] and "debug" not in opt["name"]:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank == 0:
logger.info(message)
# validation, to produce ker_map_list(fake)
if current_step % opt["train"]["val_freq"] == 0 and rank <= 0:
avg_psnr = 0.0
idx = 0
for _, val_data in enumerate(val_loader):
# LR_img, ker_map = prepro(val_data['GT'])
LR_img = val_data["LQ"]
lr_img = util.tensor2img(LR_img) # save LR image for reference
# valid Predictor
model.feed_data(LR_img, val_data["GT"])
model.test()
visuals = model.get_current_visuals()
# Save images for reference
img_name = val_data["LQ_path"][0]
img_dir = os.path.join(opt["path"]["val_images"], img_name)
# img_dir = os.path.join(opt['path']['val_images'], str(current_step), '_', str(step))
util.mkdir(img_dir)
save_lr_path = os.path.join(img_dir, "{:s}_LR.png".format(img_name))
util.save_img(lr_img, save_lr_path)
sr_img = util.tensor2img(visuals["SR"].squeeze()) # uint8
gt_img = util.tensor2img(visuals["GT"].squeeze()) # uint8
save_img_path = os.path.join(
img_dir, "{:s}_{:d}.png".format(img_name, current_step)
)
kernel = (
visuals["ker"]
.numpy()
.reshape(
opt["degradation"]["ksize"], opt["degradation"]["ksize"]
)
)
kernel = 1 / (np.max(kernel) + 1e-4) * 255 * kernel
cv2.imwrite(save_img_path, kernel)
util.save_img(sr_img, save_img_path)
# calculate PSNR
crop_size = opt["scale"]
gt_img = gt_img / 255.0
sr_img = sr_img / 255.0
cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size]
cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size]
avg_psnr += util.calculate_psnr(
cropped_sr_img * 255, cropped_gt_img * 255
)
idx += 1
avg_psnr = avg_psnr / idx
# log
logger.info("# Validation # PSNR: {:.6f}".format(avg_psnr))
logger_val = logging.getLogger("val") # validation logger
logger_val.info(
"<epoch:{:3d}, iter:{:8,d}, psnr: {:.6f}".format(
epoch, current_step, avg_psnr
)
)
# tensorboard logger
if opt["use_tb_logger"] and "debug" not in opt["name"]:
tb_logger.add_scalar("psnr", avg_psnr, current_step)
#### save models and training states
if current_step % opt["logger"]["save_checkpoint_freq"] == 0:
if rank <= 0:
logger.info("Saving models and training states.")
model.save(current_step)
model.save_training_state(epoch, current_step)
if rank <= 0:
logger.info("Saving the final model.")
model.save("latest")
logger.info("End of Predictor and Corrector training.")
tb_logger.close()
def test_AE(data_loader, model, opt, netG=None, model_=None):
print('test_AE')
import matplotlib.pyplot as plt
if netG is not None:
netG.cuda()
netG.eval()
model.eval()
batch_time = AverageMeter()
data_time = AverageMeter()
mse_losses = AverageMeter()
end_time = time.time()
output_buffer = []
previous_video_id = ''
test_results = {'results': {}}
# folder = os.path.join(opt.root_path, 'dev', 'Y')
folder = opt.result_path
if not os.path.exists(folder):
os.makedirs(fodler)
ori_clips = []
pred_clips = []
masks = []
clips = []
for i, (inputs, path) in enumerate(data_loader):
#print(i)
name = 'Y' + path[0].split('/')[-1][1:] + '.mp4'
if os.path.exists(os.path.join(folder, name)):
print(name)
continue
data_time.update(time.time() - end_time)
inputs = Variable(inputs, volatile=True)
outputs = model(inputs)
if netG is not None:
outputs = netG(outputs)
inputs = inputs[0, :, 4, :, :].cpu().data.numpy()
outputs = outputs[0, :, 0, :, :].cpu().data.numpy()
if opt.cut:
diff = outputs - inputs
tmp = (diff < 0.01) * (diff > -0.01)
#mu = tmp.mean()
#outputs = outputs-mu
outputs[tmp] = inputs[tmp]
mse_losses.update(0)
batch_time.update(time.time() - end_time)
end_time = time.time()
print('[{}/{}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'MSE {mse_losses.val:.5f} ({mse_losses.avg:.5f})\t'.format(
i + 1,
len(data_loader),
batch_time=batch_time,
data_time=data_time,
mse_losses=mse_losses))
clips.append(outputs) # 1x3x1x128x128
if opt.t_shrink:
if (i + 1) % 125 == 0: # last
clips = [tensor2img(clip, opt) for clip in clips]
final_clip = np.stack(clips)
name = 'Y' + path[0].split('/')[-1][1:] + '.mp4'
createVideoClip(final_clip, folder, name)
clips = []
print('Predicted video clip {} saving'.format(name))
else:
print('Not Implemented Error')
print('mse_losses:', mse_losses.avg)
with open(os.path.join(opt.result_path, '{}.json'.format(opt.test_subset)),
'w') as f:
json.dump(test_results, f)
parser.add_argument("-output_dir",
type=str,
default="../../../data_samples/DANv2")
args = parser.parse_args()
opt = option.parse(args.opt, is_train=False)
opt = option.dict_to_nonedict(opt)
model = create_model(opt)
if not osp.exists(args.output_dir):
os.makedirs(args.output_dir)
test_files = glob(osp.join(args.input_dir, "*png"))
for inx, path in tqdm(enumerate(test_files)):
name = path.split("/")[-1].split(".")[0]
img = cv2.imread(path)[:, :, [2, 1, 0]]
img = img.transpose(2, 0, 1)[None] / 255
img_t = torch.as_tensor(np.ascontiguousarray(img)).float()
model.feed_data(img_t)
model.test()
sr = model.fake_SR.detach().float().cpu()[0]
sr_im = util.tensor2img(sr)
save_path = osp.join(args.output_dir,
"{}_x{}.png".format(name, opt["scale"]))
cv2.imwrite(save_path, sr_im)
warp_i6.detach()), 2)
frame_pred = three_dim_model(frame_input)
frame_i3_rs = frame_i3.view(b, c, 1, h, w)
frame_target = frame_i3_rs
frame_input2 = torch.cat(
(warp_i0.detach(), warp_i1.detach(), warp_i2.detach(),
frame_pred.detach(), frame_i3_rs, warp_i4.detach(),
warp_i5.detach()), 2)
frame_pred_rf = fusion_model(
frame_input2) + frame_pred.detach()
frame_pred_rf = frame_pred_rf.view(b, c, 1, h, w)
frame_pred_rf = frame_pred_rf[:, :, :, 0:h - f_h_pad,
0:w - f_w_pad]
fusion_frame_pred = utils.tensor2img(
frame_pred_rf.view(1, c, h - f_h_pad, w - f_w_pad))
output_filename = os.path.join(output_dir, "%05d_res.png" % t)
utils.save_img(fusion_frame_pred, output_filename)
frame_input = utils.tensor2img(frame_i3.view(1, c, h, w))
output_filename = os.path.join(output_dir, "%05d_input.png" % t)
utils.save_img(frame_input, output_filename)
if len(times) > 0:
time_avg = sum(times) / len(times)
print("Average time = %f seconds (Total %d frames)" %
(time_avg, len(times)))
def train(self):
'''
Training logic for an epoch
for visualization use the following code (use batch size = 1):
category_id_to_name = {1: 'car'}
for batch_idx, dataset_dict in enumerate(self.data_loader):
if dataset_dict['idx'][0] == 10:
print(dataset_dict)
visualize(dataset_dict, category_id_to_name) #--> see this method in util
#image size: torch.Size([10, 3, 256, 256]) if batch_size = 10
'''
logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
self.data_loader.length, self.train_size))
logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
self.total_epochs, self.total_iters))
# tensorboard logger
if self.config['use_tb_logger'] and 'debug' not in self.config['name']:
version = float(torch.__version__[0:3])
if version >= 1.1: # PyTorch 1.1
from torch.utils.tensorboard import SummaryWriter
else:
logger.info(
'You are using PyTorch {}. Tensorboard will use [tensorboardX]'
.format(version))
from tensorboardX import SummaryWriter
tb_logger = SummaryWriter(log_dir='saved/tb_logger/' +
self.config['name'])
## Todo : resume capability
current_step = 0
start_epoch = 0
#### training
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(
start_epoch, current_step))
for epoch in range(start_epoch, self.total_epochs + 1):
for _, (image, targets) in enumerate(self.data_loader):
current_step += 1
if current_step > self.total_iters:
break
#### update learning rate
self.model.update_learning_rate(
current_step,
warmup_iter=self.config['train']['warmup_iter'])
#### training
self.model.feed_data(image, targets)
self.model.optimize_parameters(current_step)
#### log
if current_step % self.config['logger']['print_freq'] == 0:
logs = self.model.get_current_log()
message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
epoch, current_step,
self.model.get_current_learning_rate())
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if self.config[
'use_tb_logger'] and 'debug' not in self.config[
'name']:
tb_logger.add_scalar(k, v, current_step)
logger.info(message)
# validation
if current_step % self.config['train']['val_freq'] == 0:
self.model.test(self.valid_data_loader)
#### save models and training states
if current_step % self.config['logger'][
'save_checkpoint_freq'] == 0:
logger.info('Saving models and training states.')
self.model.save(current_step)
self.model.save_training_state(epoch, current_step)
#saving SR_images
for _, (image,
targets) in enumerate(self.valid_data_loader):
#print(image)
img_name = os.path.splitext(
os.path.basename(image['LQ_path'][0]))[0]
img_dir = os.path.join(
self.config['path']['val_images'], img_name)
mkdir(img_dir)
self.model.feed_data(image, targets)
self.model.test(self.valid_data_loader, train=False)
visuals = self.model.get_current_visuals()
sr_img = tensor2img(visuals['SR']) # uint8
gt_img = tensor2img(visuals['GT']) # uint8
lap_learned = tensor2img(
visuals['lap_learned']) # uint8
lap = tensor2img(visuals['lap']) # uint8
lap_HR = tensor2img(visuals['lap_HR']) # uint8
final_SR = tensor2img(visuals['final_SR']) # uint8
# Save SR images for reference
save_img_path = os.path.join(
img_dir,
'{:s}_{:d}_SR.png'.format(img_name, current_step))
save_img(sr_img, save_img_path)
# Save GT images for reference
save_img_path = os.path.join(
img_dir,
'{:s}_{:d}_GT.png'.format(img_name, current_step))
save_img(gt_img, save_img_path)
# Save final_SR images for reference
save_img_path = os.path.join(
img_dir, '{:s}_{:d}_final_SR.png'.format(
img_name, current_step))
save_img(final_SR, save_img_path)
# Save lap_learned images for reference
save_img_path = os.path.join(
img_dir, '{:s}_{:d}_lap_learned.png'.format(
img_name, current_step))
save_img(lap_learned, save_img_path)
# Save lap images for reference
save_img_path = os.path.join(
img_dir,
'{:s}_{:d}_lap.png'.format(img_name, current_step))
save_img(lap, save_img_path)
# Save lap images for reference
save_img_path = os.path.join(
img_dir, '{:s}_{:d}_lap_HR.png'.format(
img_name, current_step))
save_img(lap_HR, save_img_path)
logger.info('Saving the final model.')
self.model.save('latest')
logger.info('End of training.')
for test_data in test_loader:
single_img_psnr = []
single_img_ssim = []
single_img_psnr_y = []
single_img_ssim_y = []
need_GT = False if test_loader.dataset.opt["dataroot_GT"] is None else True
img_path = test_data["GT_path"][0] if need_GT else test_data["LQ_path"][0]
img_name = img_path
# img_name = os.path.splitext(os.path.basename(img_path))[0]
#### input dataset_LQ
model.feed_data(test_data["LQ"], test_data["GT"])
model.test()
visuals = model.get_current_visuals()
SR_img = visuals["Batch_SR"]
sr_img = util.tensor2img(visuals["SR"]) # uint8
suffix = opt["suffix"]
if suffix:
save_img_path = os.path.join(dataset_dir, img_name + suffix + ".png")
else:
save_img_path = os.path.join(dataset_dir, img_name + ".png")
util.save_img(sr_img, save_img_path)
if need_GT:
gt_img = util.tensor2img(visuals["GT"])
gt_img = gt_img / 255.0
sr_img = sr_img / 255.0
crop_border = opt["crop_border"] if opt["crop_border"] else opt["scale"]
if crop_border == 0:
def train_model(model, dataloader, criteria, optimizers, schedulers,
num_epochs, params):
# Note the time
since = time.time()
# Unpack parameters
writer = params['writer']
if writer is not None: board = True
txt_file = params['txt_file']
pretrained = params['model_files'][1]
pretrain = params['pretrain']
print_freq = params['print_freq']
dataset_size = params['dataset_size']
device = params['device']
batch = params['batch']
pretrain_epochs = params['pretrain_epochs']
gamma = params['gamma']
update_interval = params['update_interval']
tol = params['tol']
dl = dataloader
# Pretrain or load weights
if pretrain:
while True:
pretrained_model = pretraining(model, copy.deepcopy(dl),
criteria[0], optimizers[1],
schedulers[1], pretrain_epochs,
params)
if pretrained_model:
break
else:
for layer in model.children():
if hasattr(layer, 'reset_parameters'):
layer.reset_parameters()
model = pretrained_model
else:
try:
model.load_state_dict(torch.load(pretrained))
utils.print_both(
txt_file,
'Pretrained weights loaded from file: ' + str(pretrained))
except:
print("Couldn't load pretrained weights")
# Initialise clusters
utils.print_both(txt_file,
'\nInitializing cluster centers based on K-means')
kmeans(model, copy.deepcopy(dl), params)
utils.print_both(txt_file, '\nBegin clusters training')
# Prep variables for weights and accuracy of the best model
best_model_wts = copy.deepcopy(model.state_dict())
best_loss = 10000.0
# Initial target distribution
utils.print_both(txt_file, '\nUpdating target distribution')
output_distribution, labels, preds_prev = calculate_predictions(
model, copy.deepcopy(dl), params)
target_distribution = target(output_distribution)
nmi = utils.metrics.nmi(labels, preds_prev)
ari = utils.metrics.ari(labels, preds_prev)
acc = utils.metrics.acc(labels, preds_prev)
utils.print_both(
txt_file,
'NMI: {0:.5f}\tARI: {1:.5f}\tAcc {2:.5f}\n'.format(nmi, ari, acc))
if board:
niter = 0
writer.add_scalar('/NMI', nmi, niter)
writer.add_scalar('/ARI', ari, niter)
writer.add_scalar('/Acc', acc, niter)
update_iter = 1
finished = False
# Go through all epochs
for epoch in range(num_epochs):
utils.print_both(txt_file, 'Epoch {}/{}'.format(epoch + 1, num_epochs))
utils.print_both(txt_file, '-' * 10)
schedulers[0].step()
model.train(True) # Set model to training mode
running_loss = 0.0
running_loss_rec = 0.0
running_loss_clust = 0.0
# Keep the batch number for inter-phase statistics
batch_num = 1
img_counter = 0
# Iterate over data.
for data in dataloader:
# Get the inputs and labels
inputs, _ = data
inputs = inputs.to(device)
# Uptade target distribution, chack and print performance
if (batch_num - 1) % update_interval == 0 and not (batch_num == 1
and epoch == 0):
utils.print_both(txt_file, '\nUpdating target distribution:')
output_distribution, labels, preds = calculate_predictions(
model, dataloader, params)
target_distribution = target(output_distribution)
nmi = utils.metrics.nmi(labels, preds)
ari = utils.metrics.ari(labels, preds)
acc = utils.metrics.acc(labels, preds)
utils.print_both(
txt_file,
'NMI: {0:.5f}\tARI: {1:.5f}\tAcc {2:.5f}\t'.format(
nmi, ari, acc))
if board:
niter = update_iter
writer.add_scalar('/NMI', nmi, niter)
writer.add_scalar('/ARI', ari, niter)
writer.add_scalar('/Acc', acc, niter)
update_iter += 1
# check stop criterion
delta_label = np.sum(preds != preds_prev).astype(
np.float32) / preds.shape[0]
preds_prev = np.copy(preds)
if delta_label < tol:
utils.print_both(
txt_file, 'Label divergence ' + str(delta_label) +
'< tol ' + str(tol))
utils.print_both(
txt_file,
'Reached tolerance threshold. Stopping training.')
finished = True
break
tar_dist = target_distribution[((batch_num - 1) *
batch):(batch_num * batch), :]
tar_dist = torch.from_numpy(tar_dist).to(device)
# print(tar_dist)
# zero the parameter gradients
optimizers[0].zero_grad()
# Calculate losses and backpropagate
with torch.set_grad_enabled(True):
outputs, clusters, _ = model(inputs)
loss_rec = criteria[0](outputs, inputs)
loss_clust = gamma * criteria[1](torch.log(clusters),
tar_dist) / batch
loss = loss_rec + loss_clust
loss.backward()
optimizers[0].step()
# For keeping statistics
running_loss += loss.item() * inputs.size(0)
running_loss_rec += loss_rec.item() * inputs.size(0)
running_loss_clust += loss_rec.item() * inputs.size(0)
# Some current stats
loss_batch = loss.item()
loss_batch_rec = loss_rec.item()
loss_batch_clust = loss_clust.item()
loss_accum = running_loss / (
(batch_num - 1) * batch + inputs.size(0))
loss_accum_rec = running_loss_rec / (
(batch_num - 1) * batch + inputs.size(0))
loss_accum_clust = running_loss_clust / (
(batch_num - 1) * batch + inputs.size(0))
if batch_num % print_freq == 0:
utils.print_both(
txt_file, 'Epoch: [{0}][{1}/{2}]\t'
'Loss {3:.4f} ({4:.4f})\t'
'Loss_recovery {5:.4f} ({6:.4f})\t'
'Loss clustering {7:.4f} ({8:.4f})\t'.format(
epoch + 1, batch_num, len(dataloader), loss_batch,
loss_accum, loss_batch_rec, loss_accum_rec,
loss_batch_clust, loss_accum_clust))
if board:
niter = epoch * len(dataloader) + batch_num
writer.add_scalar('/Loss', loss_accum, niter)
writer.add_scalar('/Loss_recovery', loss_accum_rec, niter)
writer.add_scalar('/Loss_clustering', loss_accum_clust,
niter)
batch_num = batch_num + 1
# Print image to tensorboard
if batch_num == len(dataloader) and (epoch + 1) % 5:
inp = utils.tensor2img(inputs)
out = utils.tensor2img(outputs)
if board:
img = np.concatenate((inp, out), axis=1)
writer.add_image(
'Clustering/Epoch_' + str(epoch + 1).zfill(3) +
'/Sample_' + str(img_counter).zfill(2), img)
img_counter += 1
if finished: break
epoch_loss = running_loss / dataset_size
epoch_loss_rec = running_loss_rec / dataset_size
epoch_loss_clust = running_loss_clust / dataset_size
if board:
writer.add_scalar('/Loss' + '/Epoch', epoch_loss, epoch + 1)
writer.add_scalar('/Loss_rec' + '/Epoch', epoch_loss_rec,
epoch + 1)
writer.add_scalar('/Loss_clust' + '/Epoch', epoch_loss_clust,
epoch + 1)
utils.print_both(
txt_file,
'Loss: {0:.4f}\tLoss_recovery: {1:.4f}\tLoss_clustering: {2:.4f}'.
format(epoch_loss, epoch_loss_rec, epoch_loss_clust))
# If wanted to do some criterium in the future (for now useless)
if epoch_loss < best_loss or epoch_loss >= best_loss:
best_loss = epoch_loss
best_model_wts = copy.deepcopy(model.state_dict())
utils.print_both(txt_file, '')
time_elapsed = time.time() - since
utils.print_both(
txt_file,
'Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60,
time_elapsed % 60))
# load best model weights
model.load_state_dict(best_model_wts)
return model
def gen_adv():
mse = 0
num = 1
original_files = get_original_file(input_dir + val_list)
f = open('log.txt', 'w') # log
for filename, label in original_files:
img_path = input_dir + filename
print("Image: {0} ".format(img_path))
img = process_img(img_path)
Res_result, Inception_result, Mob_result = exe.run(
double_eval_program,
fetch_list=[Res_out, Inception_out, Mob_out],
feed={input_layer.name: img})
Res_result = Res_result[0]
Inception_result = Inception_result[0]
Mob_result = Mob_result[0]
r_o_label = np.argsort(Res_result)[::-1][:1][0]
i_o_label = np.argsort(Inception_result)[::-1][:1][0]
m_o_label = np.argsort(Mob_result)[::-1][:1][0]
pred_label = [r_o_label, i_o_label, m_o_label]
print("原始标签为{0}".format(label))
print("Res result: %d, Inception result: %d, Mobile result: %d" %
(r_o_label, i_o_label, m_o_label))
f.write("原始标签为{0}\n".format(label))
f.write("Res result: %d, Inception result: %d, Mobile result: %d\n" %
(r_o_label, i_o_label, m_o_label))
if r_o_label == int(label) and i_o_label == int(
label) and m_o_label == int(label):
global Res_ratio, Incep_ratio, Mob_ratio
Res_ratio = 30.0 / 43.0
Incep_ratio = 10.0 / 43.0
Mob_ratio = 3.0 / 43.0
adv_img = attack_nontarget_by_PGD(
double_adv_program,
img,
pred_label,
label,
out=[Res_out, Inception_out, Mob_out])
image_name, image_ext = filename.split('.')
##Save adversarial image(.png)
org_img = tensor2img(img)
score = calc_mse(org_img, adv_img)
#image_name += "MSE_{}".format(score)
save_adv_image(adv_img, output_dir + image_name + '.png')
mse += score
elif r_o_label == int(label) and m_o_label == int(
label): # Inception 预测错误
print("filename:{}, Inception failed!".format(filename))
Res_ratio = 0.9
Incep_ratio = 0
Mob_ratio = 0.1
adv_img = attack_nontarget_by_PGD(double_adv_program,
img, [r_o_label, 0, m_o_label],
label,
out=[Res_out, Mob_out])
image_name, image_ext = filename.split('.')
##Save adversarial image(.png)
org_img = tensor2img(img)
score = calc_mse(org_img, adv_img)
#image_name += "MSE_{}".format(score)
save_adv_image(adv_img, output_dir + image_name + '.png')
mse += score
elif r_o_label == int(label) and i_o_label == int(
label): # Mobile 预测错误
print("filename:{}, Mobile failed!".format(filename))
Res_ratio = 0.75
Incep_ratio = 0.25
Mob_ratio = 0
adv_img = attack_nontarget_by_PGD(double_adv_program,
img, [r_o_label, i_o_label, 0],
label,
out=[Res_out, Inception_out])
image_name, image_ext = filename.split('.')
##Save adversarial image(.png)
org_img = tensor2img(img)
score = calc_mse(org_img, adv_img)
# image_name += "MSE_{}".format(score)
save_adv_image(adv_img, output_dir + image_name + '.png')
mse += score
elif r_o_label == int(label): # Mobile, Inception 预测错误
print("filename:{}, Mobile failed!, Inception failed!".format(
filename))
Res_ratio = 1.0
Incep_ratio = 0.0
Mob_ratio = 0.0
adv_img = attack_nontarget_by_PGD(double_adv_program,
img, [r_o_label, 0, 0],
label,
out=[Res_out])
image_name, image_ext = filename.split('.')
##Save adversarial image(.png)
org_img = tensor2img(img)
score = calc_mse(org_img, adv_img)
# image_name += "MSE_{}".format(score)
save_adv_image(adv_img, output_dir + image_name + '.png')
mse += score
else:
print("{0}个样本已为对抗样本, name为{1}".format(num, filename))
score = 0
f.write("{0}个样本已为对抗样本, name为{1}\n".format(num, filename))
img = tensor2img(img)
image_name, image_ext = filename.split('.')
#image_name += "_un_adv_"
save_adv_image(img, output_dir + image_name + '.png')
print("this rext network weight is {0}".format(Res_ratio))
num += 1
print("the image's mse is {}".format(score))
# break
print("ADV {} files, AVG MSE: {} ".format(len(original_files),
mse / len(original_files)))
#print("ADV {} files, AVG MSE: {} ".format(len(original_files - num), mse / len(original_files - num)))
f.write("ADV {} files, AVG MSE: {} ".format(len(original_files),
mse / len(original_files)))
f.close()
dim=1)
### forward
ts = time.time()
output, lstm_state = model(inputs, lstm_state)
frame_o2 = frame_p2 + output
te = time.time()
times.append(te - ts)
## create new variable to detach from graph and avoid memory accumulation
lstm_state = utils.repackage_hidden(lstm_state)
### convert to numpy array
frame_o2 = utils.tensor2img(frame_o2)
### resize to original size
frame_o2 = cv2.resize(frame_o2, (W_orig, H_orig))
### save output frame
output_filename = os.path.join(output_dir, "%05d.jpg" % (t))
utils.save_img(frame_o2, output_filename)
## end of frame
## end of video
if len(times) > 0:
time_avg = sum(times) / len(times)
print("Average time = %f seconds (Total %d frames)" %
(time_avg, len(times)))
print(log)
if opt.visualize:
vis_log(log, vis)
vis_img(fake_A, "fake_A", vis)
vis_img(fake_B, "fake_B", vis)
vis_loss(loss_G.reshape(-1), "loss_G", vis_update, vis)
vis_loss(loss_DA.reshape(-1), "loss_DA", vis_update, vis)
vis_loss(loss_DB.reshape(-1), "loss_DB", vis_update, vis)
vis_loss(loss_cyc_A.reshape(-1), "loss_cyc_A", vis_update, vis)
vis_loss(loss_cyc_B.reshape(-1), "loss_cyc_B", vis_update, vis)
vis_update += 1
if i % 1000 == 0:
#fake_A = 0.5*(fake_A.data + 1)
#fake_B = 0.5*(fake_B.data + 1)
fake_A_img = tensor2img(fake_A[0])
fake_B_img = tensor2img(fake_B[0])
cycle_A_img = tensor2img(cycle_A[0])
cycle_B_img = tensor2img(cycle_B[0])
idt_A_img = tensor2img(idt_A[0])
idt_B_img = tensor2img(idt_B[0])
real_A_img = tensor2img(real_A[0])
real_B_img = tensor2img(real_B[0])
fake_A_img.save('%s/%s_%s_A.jpg' % (log_path, epoch, i))
fake_B_img.save('%s/%s_%s_B.jpg' % (log_path, epoch, i))
cycle_A_img.save('%s/%s_%s_cycle_A.jpg' % (log_path, epoch, i))
cycle_B_img.save('%s/%s_%s_cycle_B.jpg' % (log_path, epoch, i))
idt_A_img.save('%s/%s_%s_idt_A.jpg' % (log_path, epoch, i))
idt_B_img.save('%s/%s_%s_idt_B.jpg' % (log_path, epoch, i))
real_A_img.save('%s/%s_%s_real_A.jpg' % (log_path, epoch, i))
real_B_img.save('%s/%s_%s_real_B.jpg' % (log_path, epoch, i))
def main(use_cuda):
"""
Advbox demo which demonstrate how to use advbox.
"""
main_prog = fluid.default_main_program()
output_target = './datasets/output_image/'
if not os.path.exists(output_target):
os.makedirs(output_target)
IMG_NAME = 'img'
LABEL_NAME = 'label'
global_id = 0
img = fluid.layers.data(name=IMG_NAME,
shape=[3, 224, 224],
dtype='float32')
label = fluid.layers.data(name=LABEL_NAME, shape=[1], dtype='int64')
noise = fluid.layers.create_parameter(
name="noise",
shape=[batch_size, 3, 224, 224],
dtype='float32',
default_initializer=fluid.initializer.Constant(0.0000001))
true_image = noise + img
r_image = fluid.layers.crop(true_image,
shape=[batch_size, 1, 224, 224],
offsets=[0, 0, 0, 0],
name='r_image')
g_image = fluid.layers.crop(true_image,
shape=[batch_size, 1, 224, 224],
offsets=[0, 1, 0, 0],
name='g_image')
b_image = fluid.layers.crop(true_image,
shape=[batch_size, 1, 224, 224],
offsets=[0, 2, 0, 0],
name='b_image')
max_mean = [0.485, 0.456, 0.406]
max_std = [0.229, 0.224, 0.225]
r_max = (1 - max_mean[0]) / max_std[0]
g_max = (1 - max_mean[1]) / max_std[1]
b_max = (1 - max_mean[2]) / max_std[2]
r_min = (0 - max_mean[0]) / max_std[0]
g_min = (0 - max_mean[1]) / max_std[1]
b_min = (0 - max_mean[2]) / max_std[2]
r_image = fluid.layers.clip(x=r_image, min=r_min, max=r_max)
g_image = fluid.layers.clip(x=g_image, min=g_min, max=g_max)
b_image = fluid.layers.clip(x=b_image, min=b_min, max=b_max)
true_image = fluid.layers.concat([r_image, g_image, b_image], axis=1)
loss, outs = create_net(true_image, label)
std = fluid.layers.assign(
np.array([[[0.229]], [[0.224]], [[0.225]]]).astype('float32'))
square = fluid.layers.square(noise * std * 255.0)
# avg l2 norm
loss2 = fluid.layers.reduce_sum(square, dim=1)
loss2 = fluid.layers.sqrt(loss2)
loss2 = fluid.layers.reduce_mean(loss2)
#avg mse
# loss2 = fluid.layers.reduce_mean(square)
loss = loss + 0.005 * loss2
init_prog(main_prog)
test_prog = main_prog.clone()
lr = fluid.layers.create_global_var(shape=[1],
value=0.02,
dtype='float32',
persistable=True,
name='learning_rate_0')
opt = fluid.optimizer.Adam(learning_rate=lr)
opt.minimize(loss, parameter_list=[noise.name])
# 根据配置选择使用CPU资源还是GPU资源
place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())
test_reader = paddle.batch(test_set(), batch_size=batch_size)
load_params(exe)
fail_count = 0
for block in main_prog.blocks:
for var in block.vars.keys():
if 'learning_rate' in var:
pd_lr = fluid.global_scope().find_var(var)
print(var)
if 'beta1_pow_acc' in var:
pd_noise_beta1 = fluid.global_scope().find_var(var)
print(var)
if 'moment1' in var:
pd_noise_mom1 = fluid.global_scope().find_var(var)
print(var)
if 'beta2_pow_acc' in var:
pd_noise_beta2 = fluid.global_scope().find_var(var)
print(var)
if 'moment2' in var:
pd_noise_mom2 = fluid.global_scope().find_var(var)
print(var)
print(np.array(pd_lr.get_tensor()))
for train_id, data in enumerate(test_reader()):
images = []
labels = []
filenames = []
for i in range(batch_size):
images.append(data[i][0][0])
labels.append([data[i][1]])
filenames.append(data[i][2])
# image = data[0][0]
# label = data[0][1]
# label = np.array([[label]])
# filename = data[0][2]
images = np.array(images)
labels = np.array(labels)
for block in main_prog.blocks:
for param in block.all_parameters():
if param.name == 'noise':
pd_var = fluid.global_scope().find_var(param.name)
pd_param = pd_var.get_tensor()
print("load: {}, shape: {}".format(param.name,
param.shape))
print("Before setting the numpy array value: {}".format(
np.array(pd_param).ravel()[:5]))
noise_tensor = np.zeros(param.shape).astype('float32')
noise_tensor[:] = 1e-7
pd_param.set(noise_tensor, place)
print("After setting the numpy array value: {}".format(
np.array(pd_param).ravel()[:5]))
# pd_lr.get_tensor().set(np.array([0.02]).astype('float32'), place)
if batch_size > 1:
pd_noise_beta1.get_tensor().set(
np.array([0.9]).astype('float32'), place)
pd_noise_beta2.get_tensor().set(
np.array([0.999]).astype('float32'), place)
pd_noise_mom1.get_tensor().set(
np.zeros(shape=[batch_size, 3, 224, 224]).astype('float32'),
place)
pd_noise_mom2.get_tensor().set(
np.zeros(shape=[batch_size, 3, 224, 224]).astype('float32'),
place)
i = 0
fetch_list = [true_image, lr, loss, loss2, noise]
mean_np = np.array([[[[0.485]], [[0.456]],
[[0.406]]]]).astype('float32')
std_np = np.array([[[[0.229]], [[0.224]],
[[0.225]]]]).astype('float32')
ori_img = np.round((images * std_np + mean_np) * 255.0)
ori_img = np.clip(ori_img, 0, 255).astype('uint8')
while True:
if i == 0:
test_vars = exe.run(program=test_prog,
feed={
'img': images,
'label': labels
},
fetch_list=outs)
for m in range(batch_size):
str = 'First step test network,id:{},'.format(global_id +
1)
global_id += 1
adv_labels = []
for j in range(len(outs)):
o = test_vars[j][m]
adv_label = np.argmax(o)
adv_labels.append(adv_label)
str += 'adv{}:%d,'.format(j + 1)
print(str % (*adv_labels, ))
train_vars = exe.run(program=fluid.default_main_program(),
feed={
'img': images,
'label': labels
},
fetch_list=fetch_list)
n = train_vars[-1]
l2 = train_vars[-2]
l1 = train_vars[-3]
lr1 = train_vars[-4]
tr_img = train_vars[-5]
adv_img = n + images
adv_img = np.round((adv_img * std_np + mean_np) * 255.0)
adv_img = np.clip(adv_img, 0, 255).astype('uint8')
diff = adv_img.astype('float32') - ori_img.astype('float32')
avg_mse = diff * diff
# avg l2 norm
l2_norm = np.sum(avg_mse, axis=1)
l2_norm = np.sqrt(l2_norm)
l2_norm = np.mean(l2_norm, axis=(1, 2))
# avg mse
avg_mse = np.mean(avg_mse, axis=(1, 2, 3))
test_vars = exe.run(program=test_prog,
feed={
'img': images,
'label': labels
},
fetch_list=outs)
successful = batch_size * len(outs)
for m in range(batch_size):
str = 'batch:%d,id:{},lr:%f,loss1:%f,loss2:%f,avg_mse:%f,l2_norm:%f,'.format(
train_id * batch_size + m + 1)
adv_labels = []
for j in range(len(outs)):
o = test_vars[j][m]
adv_label = np.argmax(o)
adv_labels.append(adv_label)
str += 'adv{}:%d,'.format(j + 1)
print(str %
(i, lr1, l1, l2, avg_mse[m], l2_norm[m], *adv_labels))
for adv_label in adv_labels:
if adv_label == labels[m]:
successful -= 1
i += 1
if (successful >= batch_size * len(outs) - 1
and np.mean(l2_norm) < 1.0) or i == 3000:
if successful >= batch_size * len(outs) - 1:
print('attack successful')
else:
print('attack failed')
fail_count += 1
break
print("failed:%d" % (fail_count, ))
adv_img = adv_img.astype('float32') / 255.0
adv_img = adv_img - mean_np
adv_img = adv_img / std_np
for m in range(batch_size):
adv_image = tensor2img(adv_img[m][np.newaxis, :, :, :])
ori_image = tensor2img(images[m][np.newaxis, :, :, :])
print('id:{},mse:{}'.format(train_id * batch_size + m + 1,
call_avg_mse_np(adv_image, ori_image)))
save_adv_image(
adv_image,
os.path.join(output_target,
filenames[m].split('.')[0] + '.png'))
print("attack over ,failed:%d" % (fail_count, ))
def pretraining(model, dataloader, criterion, optimizer, scheduler, num_epochs,
params):
# Note the time
since = time.time()
# Unpack parameters
writer = params['writer']
if writer is not None: board = True
txt_file = params['txt_file']
pretrained = params['model_files'][1]
print_freq = params['print_freq']
dataset_size = params['dataset_size']
device = params['device']
batch = params['batch']
# Prep variables for weights and accuracy of the best model
best_model_wts = copy.deepcopy(model.state_dict())
best_loss = 10000.0
# Go through all epochs
for epoch in range(num_epochs):
utils.print_both(
txt_file,
'Pretraining:\tEpoch {}/{}'.format(epoch + 1, num_epochs))
utils.print_both(txt_file, '-' * 10)
scheduler.step()
model.train(True) # Set model to training mode
running_loss = 0.0
# Keep the batch number for inter-phase statistics
batch_num = 1
# Images to show
img_counter = 0
# Iterate over data.
for data in dataloader:
# Get the inputs and labels
inputs, _ = data
inputs = inputs.to(device)
# zero the parameter gradients
optimizer.zero_grad()
with torch.set_grad_enabled(True):
outputs, _, _ = model(inputs)
loss = criterion(outputs, inputs)
loss.backward()
optimizer.step()
# For keeping statistics
running_loss += loss.item() * inputs.size(0)
# Some current stats
loss_batch = loss.item()
loss_accum = running_loss / (
(batch_num - 1) * batch + inputs.size(0))
if batch_num % print_freq == 0:
utils.print_both(
txt_file, 'Pretraining:\tEpoch: [{0}][{1}/{2}]\t'
'Loss {3:.4f} ({4:.4f})\t'.format(epoch + 1, batch_num,
len(dataloader),
loss_batch, loss_accum))
if board:
niter = epoch * len(dataloader) + batch_num
writer.add_scalar('Pretraining/Loss', loss_accum, niter)
batch_num = batch_num + 1
if batch_num in [
len(dataloader),
len(dataloader) // 2,
len(dataloader) // 4, 3 * len(dataloader) // 4
]:
inp = utils.tensor2img(inputs)
out = utils.tensor2img(outputs)
if board:
img = np.concatenate((inp, out), axis=1)
writer.add_image(
'Pretraining/Epoch_' + str(epoch + 1).zfill(3) +
'/Sample_' + str(img_counter).zfill(2), img)
img_counter += 1
epoch_loss = running_loss / dataset_size
if epoch == 0: first_loss = epoch_loss
if epoch == 4 and epoch_loss / first_loss > 1:
utils.print_both(
txt_file,
"\nLoss not converging, starting pretraining again\n")
return False
if board:
writer.add_scalar('Pretraining/Loss' + '/Epoch', epoch_loss,
epoch + 1)
utils.print_both(txt_file,
'Pretraining:\t Loss: {:.4f}'.format(epoch_loss))
# If wanted to add some criterium in the future
if epoch_loss < best_loss or epoch_loss >= best_loss:
best_loss = epoch_loss
best_model_wts = copy.deepcopy(model.state_dict())
utils.print_both(txt_file, '')
time_elapsed = time.time() - since
utils.print_both(
txt_file, 'Pretraining complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
# load best model weights
model.load_state_dict(best_model_wts)
model.pretrained = True
torch.save(model.state_dict(), pretrained)
return model
