File size: 1,940 Bytes
998bb30 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import torch
import torch.nn as nn
class Trigger(nn.Module):
def __init__(self, size: int = 32, epsilon: float=16/255, transparency: float = 1.) -> None:
super().__init__()
self.size = size
self.epsilon = epsilon
self.mask = nn.Parameter(torch.rand(size, size,device=torch.device('cuda')),requires_grad=True)
self.transparency = transparency
self.trigger = nn.Parameter(torch.rand(3, size, size,device=torch.device('cuda')) * 4 - 2,requires_grad=True)
# self.trigger = nn.Parameter(torch.rand(3, size, size,device=torch.device('cuda')),requires_grad=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x = torch.min(torch.max(tri(x), x - epsilon), x + epsilon)
if len(x.shape) == 4:
return self.transparency * self.mask * self.trigger.repeat(len(x), 1, 1, 1) + (1 - self.mask * self.transparency) * x
# return torch.clamp(self.transparency * torch.clamp(self.mask, 0, 1) * torch.clamp(self.trigger.repeat(len(x), 1, 1, 1), -self.epsilon,
# self.epsilon) + (
# 1 - torch.clamp(self.mask, 0, 1) * self.transparency) * x, -1, 1)
else:
return self.transparency * self.mask * self.trigger + (1 - self.mask * self.transparency) * x
# return torch.clamp(self.transparency * torch.clamp(self.mask, 0, 1) * torch.clamp(self.trigger, -self.epsilon, self.epsilon) + (1 - torch.clamp(self.mask, 0, 1) * self.transparency) * x, -1, 1)
class UAP(nn.Module):
def __init__(self, size: int = 32) -> None:
super().__init__()
self.size = size
self.perturbation = nn.Parameter(torch.zeros(3, size, size,device=torch.device('cuda')),requires_grad=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return x + self.perturbation
|