226 lines
9.6 KiB
Python
226 lines
9.6 KiB
Python
"""
|
||
Training code for Adversarial patch training
|
||
|
||
|
||
"""
|
||
|
||
import PIL
|
||
from torch.utils.tensorboard import SummaryWriter
|
||
|
||
# import load_data
|
||
from tqdm import tqdm
|
||
|
||
from load_data import * # 可能导致多次导入问题？
|
||
import gc
|
||
import matplotlib.pyplot as plt
|
||
from torch import autograd
|
||
from torchvision import transforms
|
||
|
||
import subprocess
|
||
|
||
import patch_config
|
||
import sys
|
||
import time
|
||
|
||
from yolo import YOLO
|
||
|
||
|
||
class PatchTrainer(object):
|
||
def __init__(self, mode):
|
||
self.config = patch_config.patch_configs[mode]() # 获取对应的配置类
|
||
|
||
# self.darknet_model = Darknet(self.config.cfgfile) # 加载yolo模型
|
||
# self.darknet_model.load_weights(self.config.weightfile) # 默认 YOLOv2 MS COCO weights， person编号是0
|
||
self.darknet_model = YOLO().net
|
||
self.darknet_model = self.darknet_model.eval().cuda() # TODO: Why eval?
|
||
self.patch_applier = PatchApplier().cuda() # 对图像应用对抗补丁
|
||
self.patch_transformer = PatchTransformer().cuda() # 变换补丁到指定大小并产生抖动
|
||
# self.prob_extractor = MaxProbExtractor(0, 80, self.config).cuda() # 提取最大类别概率
|
||
self.prob_extractor = MaxProbExtractor(0, 1, self.config).cuda() # 提取最大类别概率
|
||
self.nps_calculator = NPSCalculator(self.config.printfile, self.config.patch_size).cuda() # 不可打印分数
|
||
self.total_variation = TotalVariation().cuda() # 计算补丁的所有变化程度
|
||
|
||
self.writer = self.init_tensorboard(mode)
|
||
|
||
def init_tensorboard(self, name=None):
|
||
subprocess.Popen(['tensorboard', '--logdir=runs'])
|
||
if name is not None:
|
||
time_str = time.strftime("%Y%m%d-%H%M%S")
|
||
return SummaryWriter(f'runs/{time_str}_{name}')
|
||
else:
|
||
return SummaryWriter()
|
||
|
||
def train(self):
|
||
"""
|
||
Optimize a patch to generate an adversarial example.
|
||
:return: Nothing
|
||
"""
|
||
|
||
img_size = self.darknet_model.height # 416
|
||
# print('batch_size:',batch_size)
|
||
batch_size = self.config.batch_size # 8
|
||
n_epochs = 200
|
||
# n_epochs = 5
|
||
# max_lab = 20 # label的最大长度
|
||
max_lab = 8
|
||
|
||
time_str = time.strftime("%Y%m%d-%H%M%S")
|
||
|
||
# Generate stating point
|
||
# adv_patch_cpu = self.generate_patch("gray") # 生成一个灰图，初始化为0.5
|
||
adv_patch_cpu = self.read_image("saved_patches/patchnew0.jpg")
|
||
|
||
adv_patch_cpu.requires_grad_(True)
|
||
|
||
train_loader = torch.utils.data.DataLoader(
|
||
InriaDataset(self.config.img_dir, self.config.lab_dir, max_lab, img_size,
|
||
shuffle=True),
|
||
batch_size=batch_size,
|
||
shuffle=True,
|
||
num_workers=0) # 与 from load_data import * 搭配导致多少导入？
|
||
self.epoch_length = len(train_loader)
|
||
print(f'One epoch is {len(train_loader)}')
|
||
|
||
optimizer = optim.Adam([adv_patch_cpu], lr=self.config.start_learning_rate, amsgrad=True) # 更新的是那个补丁
|
||
scheduler = self.config.scheduler_factory(optimizer) # ICLR-2018年最佳论文提出的Adam改进版Amsgrad
|
||
|
||
et0 = time.time()
|
||
for epoch in range(n_epochs):
|
||
ep_det_loss = 0
|
||
ep_nps_loss = 0
|
||
ep_tv_loss = 0
|
||
ep_loss = 0
|
||
bt0 = time.time()
|
||
for i_batch, (img_batch, lab_batch) in tqdm(enumerate(train_loader), desc=f'Running epoch {epoch}',
|
||
total=self.epoch_length):
|
||
with autograd.detect_anomaly(): # 1.运行前向时开启异常检测功能，则在反向时会打印引起反向失败的前向操作堆栈 2.反向计算出现“nan”时引发异常
|
||
img_batch = img_batch.cuda() # 8, 3, 416, 416
|
||
lab_batch = lab_batch.cuda() # 8, 14, 5 为什么要把人数的标签补到14?
|
||
# print('TRAINING EPOCH %i, BATCH %i'%(epoch, i_batch))
|
||
adv_patch = adv_patch_cpu.cuda() # 3, 300, 300
|
||
adv_batch_t = self.patch_transformer(adv_patch, lab_batch, img_size, do_rotate=True, rand_loc=False)
|
||
p_img_batch = self.patch_applier(img_batch, adv_batch_t)
|
||
p_img_batch = F.interpolate(p_img_batch,
|
||
(self.darknet_model.height, self.darknet_model.width)) # 确保和图片大小一致
|
||
|
||
# print('++++++++++++p_img_batch:+++++++++++++',p_img_batch.shape)
|
||
img = p_img_batch[1, :, :, ]
|
||
img = transforms.ToPILImage()(img.detach().cpu())
|
||
# img.show()
|
||
|
||
outputs = self.darknet_model(p_img_batch) # 输入8，3，416，416 输出8，425， 13， 13 ,其中425是5*(5+80)
|
||
max_prob = 0
|
||
nps = 0
|
||
tv = 0
|
||
for l in range(len(outputs)): # 三组不同分辨率大小的输出特征分别计算
|
||
output = outputs[l]
|
||
max_prob += self.prob_extractor(output)
|
||
nps += self.nps_calculator(adv_patch)
|
||
tv += self.total_variation(adv_patch)
|
||
|
||
nps_loss = nps * 0.01
|
||
tv_loss = tv * 2.5
|
||
det_loss = torch.mean(max_prob) # 把人的置值度当成损失
|
||
loss = det_loss + nps_loss + torch.max(tv_loss, torch.tensor(0.1).cuda())
|
||
|
||
ep_det_loss += det_loss.detach().cpu().numpy()
|
||
ep_nps_loss += nps_loss.detach().cpu().numpy()
|
||
ep_tv_loss += tv_loss.detach().cpu().numpy()
|
||
ep_loss += loss
|
||
|
||
loss.backward()
|
||
optimizer.step()
|
||
optimizer.zero_grad()
|
||
adv_patch_cpu.data.clamp_(0, 1) # keep patch in image range
|
||
|
||
bt1 = time.time()
|
||
if i_batch % 5 == 0:
|
||
iteration = self.epoch_length * epoch + i_batch
|
||
|
||
self.writer.add_scalar('total_loss', loss.detach().cpu().numpy(), iteration)
|
||
self.writer.add_scalar('loss/det_loss', det_loss.detach().cpu().numpy(), iteration)
|
||
self.writer.add_scalar('loss/nps_loss', nps_loss.detach().cpu().numpy(), iteration)
|
||
self.writer.add_scalar('loss/tv_loss', tv_loss.detach().cpu().numpy(), iteration)
|
||
self.writer.add_scalar('misc/epoch', epoch, iteration)
|
||
self.writer.add_scalar('misc/learning_rate', optimizer.param_groups[0]["lr"], iteration)
|
||
|
||
self.writer.add_image('patch', adv_patch_cpu, iteration)
|
||
if i_batch + 1 >= len(train_loader):
|
||
print('\n')
|
||
else:
|
||
del adv_batch_t, output, max_prob, det_loss, p_img_batch, nps_loss, tv_loss, loss
|
||
torch.cuda.empty_cache()
|
||
bt0 = time.time()
|
||
et1 = time.time()
|
||
ep_det_loss = ep_det_loss / len(train_loader)
|
||
ep_nps_loss = ep_nps_loss / len(train_loader)
|
||
ep_tv_loss = ep_tv_loss / len(train_loader)
|
||
ep_loss = ep_loss / len(train_loader)
|
||
|
||
# im = transforms.ToPILImage('RGB')(adv_patch_cpu)
|
||
# plt.imshow(im)
|
||
# plt.savefig(f'pics/{time_str}_{self.config.patch_name}_{epoch}.png')
|
||
|
||
scheduler.step(ep_loss)
|
||
if True:
|
||
print(' EPOCH NR: ', epoch),
|
||
print('EPOCH LOSS: ', ep_loss)
|
||
print(' DET LOSS: ', ep_det_loss)
|
||
print(' NPS LOSS: ', ep_nps_loss)
|
||
print(' TV LOSS: ', ep_tv_loss)
|
||
print('EPOCH TIME: ', et1 - et0)
|
||
# im = transforms.ToPILImage('RGB')(adv_patch_cpu)
|
||
# plt.imshow(im)
|
||
# plt.show()
|
||
# im.save("saved_patches/patchnew1.jpg")
|
||
im = transforms.ToPILImage('RGB')(adv_patch_cpu)
|
||
if epoch >= 3:
|
||
im.save(f"saved_patches/patchnew1_t1_{epoch}_{time_str}.jpg")
|
||
del adv_batch_t, output, max_prob, det_loss, p_img_batch, nps_loss, tv_loss, loss
|
||
torch.cuda.empty_cache()
|
||
et0 = time.time()
|
||
|
||
def generate_patch(self, type):
|
||
"""
|
||
Generate a random patch as a starting point for optimization.
|
||
|
||
:param type: Can be 'gray' or 'random'. Whether or not generate a gray or a random patch.
|
||
:return:
|
||
"""
|
||
if type == 'gray':
|
||
adv_patch_cpu = torch.full((3, self.config.patch_size, self.config.patch_size), 0.5)
|
||
elif type == 'random':
|
||
adv_patch_cpu = torch.rand((3, self.config.patch_size, self.config.patch_size))
|
||
|
||
return adv_patch_cpu
|
||
|
||
def read_image(self, path):
|
||
"""
|
||
Read an input image to be used as a patch
|
||
|
||
:param path: Path to the image to be read.
|
||
:return: Returns the transformed patch as a pytorch Tensor.
|
||
"""
|
||
patch_img = Image.open(path).convert('RGB')
|
||
tf = transforms.Resize((self.config.patch_size, self.config.patch_size))
|
||
patch_img = tf(patch_img)
|
||
tf = transforms.ToTensor()
|
||
|
||
adv_patch_cpu = tf(patch_img)
|
||
return adv_patch_cpu
|
||
|
||
|
||
def main():
|
||
if len(sys.argv) != 2:
|
||
print('You need to supply (only) a configuration mode.')
|
||
print('Possible modes are:')
|
||
print(patch_config.patch_configs) # 一般传入paper_obj
|
||
|
||
# print('sys.argv:',sys.argv)
|
||
trainer = PatchTrainer(sys.argv[1])
|
||
trainer.train()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|