|
|
import unittest
|
|
|
|
|
|
import torch
|
|
|
|
|
|
from gat.models.attack.aim_attack import AIMAttack
|
|
|
from gat.models.attack.cda_attack import CDAAttack
|
|
|
from gat.models.attack.generator.aim import AIMGenerator
|
|
|
from gat.models.attack.generator.cda import CDAGenerator
|
|
|
from gat.models.attack.loss.logits import ContrastiveLoss
|
|
|
|
|
|
|
|
|
class TestGenerator(unittest.TestCase):
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_cda(self):
|
|
|
generator = CDAGenerator()
|
|
|
inputs = torch.rand(1, 3, 224, 224)
|
|
|
outputs = generator(inputs)
|
|
|
self.assertEqual(outputs.shape, (1, 3, 224, 224))
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_aim(self):
|
|
|
generator = AIMGenerator()
|
|
|
inputs = (torch.rand(1, 3, 224, 224), torch.rand(1, 3, 224, 224))
|
|
|
outputs = generator(*inputs)
|
|
|
self.assertEqual(outputs.shape, (1, 3, 224, 224))
|
|
|
|
|
|
|
|
|
class TestContrastiveLoss(unittest.TestCase):
|
|
|
|
|
|
def setUp(self):
|
|
|
self.margin = 1.0
|
|
|
self.loss_fn = ContrastiveLoss(margin=self.margin)
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_forward_shape(self):
|
|
|
anchors = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
|
positives = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
|
negatives = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
|
|
|
|
|
|
loss = self.loss_fn(anchors, negatives, positives)
|
|
|
self.assertEqual(loss.shape, torch.Size([]))
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_loss_outputs(self):
|
|
|
anchors = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
|
positives = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
|
negatives = torch.tensor([[2.0, 0.0], [0.0, 2.0]])
|
|
|
|
|
|
loss = self.loss_fn(anchors, negatives, positives)
|
|
|
self.assertAlmostEqual(loss.item(), 0.25)
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_non_zero_loss(self):
|
|
|
anchors = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
|
positives = torch.tensor([[0.0, 1.0], [1.0, 0.0]])
|
|
|
negatives = torch.tensor([[2.0, 0.0], [0.0, 2.0]])
|
|
|
|
|
|
loss = self.loss_fn(anchors, negatives, positives)
|
|
|
self.assertAlmostEqual(loss.item(), 0.75)
|
|
|
|
|
|
|
|
|
class TestCDAAttack(unittest.TestCase):
|
|
|
|
|
|
def setUp(self):
|
|
|
self.device = torch.device('cpu')
|
|
|
self.epsilon = 16. / 255.
|
|
|
self.attack = CDAAttack(self.device, self.epsilon)
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_outputs_shape(self):
|
|
|
x_nat = torch.rand(1, 3, 224, 224)
|
|
|
x_adv = self.attack(x_nat)
|
|
|
self.assertEqual(x_adv.shape, (1, 3, 224, 224))
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_outputs_bound(self):
|
|
|
x_nat = torch.rand(1, 3, 224, 224)
|
|
|
x_adv = self.attack(x_nat)
|
|
|
self.assertTrue((x_adv >= 0.0).all())
|
|
|
self.assertTrue((x_adv <= 1.0).all())
|
|
|
self.assertTrue((x_adv - x_nat).abs().max() <= self.epsilon)
|
|
|
|
|
|
|
|
|
class TestAIMAttack(unittest.TestCase):
|
|
|
|
|
|
def setUp(self):
|
|
|
self.device = torch.device('cpu')
|
|
|
self.epsilon = 16. / 255.
|
|
|
self.attack = AIMAttack(self.device, self.epsilon)
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_outputs_shape(self):
|
|
|
x_nat = torch.rand(1, 3, 224, 224)
|
|
|
x_guid = torch.rand(1, 3, 224, 224)
|
|
|
x_adv = self.attack(x_nat, x_guid)
|
|
|
self.assertEqual(x_adv.shape, (1, 3, 224, 224))
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_outputs_bound(self):
|
|
|
x_nat = torch.rand(1, 3, 224, 224)
|
|
|
x_guid = torch.rand(1, 3, 224, 224)
|
|
|
x_adv = self.attack(x_nat, x_guid)
|
|
|
self.assertTrue((x_adv >= 0.0).all())
|
|
|
self.assertTrue((x_adv <= 1.0).all())
|
|
|
self.assertTrue((x_adv - x_nat).abs().max() <= self.epsilon)
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
unittest.main()
|
|
|
|
