def update_D(self, train_iter, is_print=True, loss_weight=1):
gen = self.get_optimizer('gen').target
dis_opt = self.get_optimizer('dis')
dis = dis_opt.target
label, real_g, condition = self.real_batch(train_iter)
g_out = gen([label, condition])
fake_g = g_out
real_d = dis([label, real_g, condition])[0]
fake_d = dis([label, fake_g, condition])[0]
d_loss = dis_loss(self.opt,
real_d,
fake_d,
real_g,
fake_g,
observer=dis if is_print else None)
d_loss *= loss_weight
fake_g.unchain_backward()
dis.cleargrads()
d_loss.backward()
dis_opt.update()
def update_core(self):
g_opt = self.get_optimizer('gen')
d_opt = self.get_optimizer('dis')
#predict
x, real_g = self.real_batch('main')
fake_g = g_opt.target(x)
self.img4save = [
cuda.to_cpu(x.array[0]),
cuda.to_cpu(real_g.array[0]),
cuda.to_cpu(fake_g.array[0])
]
real_d = d_opt.target(real_g)
fake_d = d_opt.target(fake_g)
#generator loss
g_loss = gen_loss(self.opt,
fake_d,
real_g,
fake_g,
observer=g_opt.target)
g_opt.target.cleargrads()
g_loss.backward()
g_opt.update()
#discriminator loss
x.unchain_backward()
fake_g.unchain_backward()
d_loss = dis_loss(self.opt, real_d, fake_d, observer=d_opt.target)
d_opt.target.cleargrads()
d_loss.backward()
d_opt.update()
if self.learn_from_unlabel:
#predict
unlabel_x, _ = self.real_batch('semi')
unlabel_g = g_opt.target(unlabel_x)
unlabel_d = d_opt.target(unlabel_g)
self.semi_img4save = [
cuda.to_cpu(unlabel_x.array[0]), None,
cuda.to_cpu(unlabel_g.array[0])
]
#semi-supervised loss
semi_loss = gen_semi_loss(self.opt,
unlabel_d,
unlabel_g,
observer=g_opt.target)
g_opt.target.cleargrads()
semi_loss.backward()
g_opt.update()
def test_patch(self):
angle_success = torch.zeros(self.num_angles)
total_loss = 0.0
n = 0.0
for mesh in self.mesh_dataset:
mesh = mesh.extend(self.num_angles)
#mesh_texture = mesh.textures.maps_padded()
for bg in self.test_bg_dataset:
#mesh_texture[:, 575:675, 475:575, :] = self.patch[None]
texture_image = mesh.textures.atlas_padded()
mesh.textures._atlas_padded[0, self.idx, :, :, :] = self.patch
mesh.textures.atlas = mesh.textures._atlas_padded
mesh.textures._atlas_list = None
#images = self.render_mesh_on_bg(mesh, bg)
rand_translation = torch.randint(-100, 100, (2, ))
images = self.render_mesh_on_bg(
mesh,
bg,
x_translation=rand_translation[0].item(),
y_translation=rand_translation[1].item())
reshape_img = images[:, :, :, :3].permute(0, 3, 1, 2)
reshape_img = reshape_img.to(self.device)
output = self.dnet(reshape_img)
d_loss = dis_loss(output, self.dnet.num_classes,
self.dnet.anchors, self.dnet.num_anchors, 0)
for angle in range(self.num_angles):
acc_loss = calc_acc(output[angle], self.dnet.num_classes,
self.dnet.num_anchors, 0)
angle_success[angle] += acc_loss.item()
tv = self.total_variation(self.patch)
tv_loss = tv * 2.5
loss = d_loss + torch.sum(
torch.max(tv_loss,
torch.tensor(0.1).to(self.device)))
total_loss += loss.item()
n += 1.0
unseen_success_rate = angle_success.mean() / len(self.test_bg_dataset)
print('Unseen bg success rate: ', unseen_success_rate.item())
def attack(self):
train_bgs = DataLoader(
self.bg_dataset,
batch_size=self.config.batch_size,
shuffle=True,
num_workers=1)
if self.patch is None or self.idx is None:
self.initialize_patch()
mesh = self.mesh_dataset.meshes[0]
total_variation = TotalVariation_3d(mesh, self.idx).to(self.device)
optimizer = torch.optim.SGD([self.patch], lr=1e-1, momentum=0.9)
for epoch in range(self.config.epochs):
ep_loss = 0.0
ep_acc = 0.0
n = 0.0
for mesh in self.mesh_dataset:
# Copy mesh for each camera angle
mesh = mesh.extend(self.num_angles_train)
for bg_batch in train_bgs:
bg_batch = bg_batch.to(self.device)
# To enable random camera distance training, uncomment this line:
# self.change_cameras('train', camera_dist=random.uniform(1.4, 3.0))
optimizer.zero_grad()
texture_image = mesh.textures.atlas_padded()
# Random patch augmentation
contrast = torch.FloatTensor(1).uniform_(self.min_contrast, self.max_contrast).to(self.device)
brightness = torch.FloatTensor(1).uniform_(self.min_brightness, self.max_brightness).to(self.device)
noise = torch.FloatTensor(self.patch.shape).uniform_(-1, 1) * self.noise_factor
noise = noise.to(self.device)
augmented_patch = (self.patch * contrast) + brightness + noise
# Clamp patch to avoid PyTorch3D issues
clamped_patch = augmented_patch.clone().clamp(min=1e-6, max=0.99999)
mesh.textures._atlas_padded[:,self.idx,:,:,:] = clamped_patch
mesh.textures.atlas = mesh.textures._atlas_padded
mesh.textures._atlas_list = None
# Render mesh onto background image
rand_translation = torch.randint(
-self.config.rand_translation,
self.config.rand_translation,
(2,)
)
images = self.render_mesh_on_bg_batch(mesh, bg_batch, self.num_angles_train, x_translation=rand_translation[0].item(),
y_translation=rand_translation[1].item())
reshape_img = images[:,:,:,:3].permute(0, 3, 1, 2)
reshape_img = reshape_img.to(self.device)
# Run detection model on images
output = self.dnet(reshape_img)
d_loss = dis_loss(output, self.dnet.num_classes, self.dnet.anchors, self.dnet.num_anchors, 0)
acc_loss = calc_acc(output, self.dnet.num_classes, self.dnet.num_anchors, 0)
tv = total_variation(self.patch)
tv_loss = tv * 2.5
loss = d_loss + tv_loss
ep_loss += loss.item()
ep_acc += acc_loss.item()
n += bg_batch.shape[0]
loss.backward(retain_graph=True)
optimizer.step()
# Save image and print performance statistics
patch_save = self.patch.cpu().detach().clone()
idx_save = self.idx.cpu().detach().clone()
torch.save(patch_save, 'patch_save.pt')
torch.save(idx_save, 'idx_save.pt')
save_image(reshape_img[0].cpu().detach(), "TEST_RENDER.png")
print('epoch={} loss={} success_rate={}'.format(
epoch,
(ep_loss / n),
(ep_acc / n) / self.num_angles_train)
)
if epoch % 5 == 0:
self.test_patch()
self.change_cameras('train')
def attack(self):
train_bgs = DataLoader(self.bg_dataset,
batch_size=self.config.batch_size,
shuffle=True,
num_workers=1)
mesh = self.mesh_dataset.meshes[0]
print(self.patch.shape)
total_variation = TotalVariation().cuda()
optimizer = torch.optim.SGD([self.patch], lr=1.0, momentum=0.9)
for epoch in range(self.config.epochs):
ep_loss = 0.0
ep_acc = 0.0
n = 0.0
for mesh in self.mesh_dataset:
# Copy mesh for each camera angle
mesh = mesh.extend(self.num_angles)
#mesh_texture = mesh.textures.maps_padded()
#c = 0
for bg_batch in train_bgs:
#c = c+1
#print('iter'+ str(c))
bg_batch = bg_batch.to(self.device)
optimizer.zero_grad()
# Apply patch to mesh texture (hard coded for now)
#mesh_texture[:, 575:675, 475:575, :] = self.patch[None]
texture_image = mesh.textures.atlas_padded()
mesh.textures._atlas_padded[0,
self.idx, :, :, :] = self.patch
mesh.textures.atlas = mesh.textures._atlas_padded
mesh.textures._atlas_list = None
# Render mesh onto background image
# images = self.render_mesh_on_bg(mesh, bg)
#images = self.render_mesh_on_bg_batch(mesh, bg_batch)
rand_translation = torch.randint(-100, 100, (2, ))
images = self.render_mesh_on_bg_batch(
mesh,
bg_batch,
x_translation=rand_translation[0].item(),
y_translation=rand_translation[1].item())
# print('images: ', images.shape)
reshape_img = images[:, :, :, :3].permute(0, 3, 1, 2)
reshape_img = reshape_img.to(self.device)
# Run detection model on images
output = self.dnet(reshape_img)
# Compute losses:
d_loss = dis_loss(output, self.dnet.num_classes,
self.dnet.anchors, self.dnet.num_anchors,
0)
acc_loss = calc_acc(output, self.dnet.num_classes,
self.dnet.num_anchors, 0)
tv = self.total_variation(self.patch)
tv_loss = tv * 2.5
loss = d_loss + torch.sum(
torch.max(tv_loss,
torch.tensor(0.1).to(self.device)))
ep_loss += loss.item()
ep_acc += acc_loss.item()
n += bg_batch.shape[0]
#TODO: Remove Retain Graph
loss.backward(retain_graph=True)
optimizer.step()
# Save image and print performance statistics
patch_save = self.patch.cpu().detach().clone()
idx_save = self.idx.cpu().detach().clone()
# torch.save(patch_save, 'patch_save.pt')
# torch.save(idx_save, 'idx_save.pt')
#save_image(self.patch.cpu().detach().permute(2, 0, 1), self.config.output + '_{}.png'.format(epoch))
print('epoch={} loss={} success_rate={}'.format(
epoch, (ep_loss / n), (ep_acc / n) / self.num_angles))
self.test_patch()
#TODO: Pass the variable value
if epoch % 10 == 0:
self.test_patch_faster_rcnn(
path_to_checkpoint="faster_rcnn/model-180000.pth",
dataset_name="coco2017",
backbone_name="resnet101",
prob_thresh=0.6)