import torch import torch.nn as nn import torch.nn.functional as F import numpy as np from methods.tool_func import * def preprocessing(x_fea): # x_fea: [B, 197, 384] --> x_cls_fea [B, 1, 384], x_patch_fea [B, 384, 14, 14] B, num, dim = x_fea.size()[0], x_fea.size()[1], x_fea.size()[2] x_cls_fea = x_fea[:, :1, :] x_patch_fea = x_fea[:,1:, :] x_patch_fea = x_patch_fea.contiguous().view(B,dim,num-1).view(B, dim, 14, 14) return x_cls_fea, x_patch_fea def postprocessing(x_cls_fea, x_patch_fea): # x_cls_fea [B, 1, 384], x_patch_fea [B, 384, 14, 14] --> x_fea: [B, 197, 384] B, num, dim = x_patch_fea.size()[0], x_patch_fea.size()[2]*x_patch_fea.size()[3]+1, x_patch_fea.size()[1] x_patch_fea = x_patch_fea.contiguous().view(B,dim,num-1).view(B,num-1,dim) x_fea = torch.cat((x_cls_fea, x_patch_fea), 1) return x_fea def changeNewAdvStyle_ViT(vit_fea, new_styleAug_mean, new_styleAug_std, p_thred): if(new_styleAug_mean=='None'): return vit_fea #final p = np.random.uniform() if( p < p_thred): return vit_fea cls_fea, input_fea = preprocessing(vit_fea) feat_size = input_fea.size() ori_style_mean, ori_style_std = calc_mean_std(input_fea) #print('ori mean:', ori_style_mean.mean(), 'ori std:', ori_style_std.mean()) #print('adv mean:', new_styleAug_mean.mean(), 'adv std:', new_styleAug_std.mean()) #print('mean diff:', new_styleAug_mean.mean() - ori_style_mean.mean(), 'std diff:', new_styleAug_std.mean() - ori_style_std.mean()) normalized_fea = (input_fea - ori_style_mean.expand(feat_size)) / ori_style_std.expand(feat_size) styleAug_fea = normalized_fea * new_styleAug_std.expand(feat_size) + new_styleAug_mean.expand(feat_size) styleAug_fea_vit = postprocessing(cls_fea, styleAug_fea) return styleAug_fea_vit class ProtoNet(nn.Module): def __init__(self, backbone): super().__init__() # bias & scale of cosine classifier self.bias = nn.Parameter(torch.FloatTensor(1).fill_(0), requires_grad=True) self.scale_cls = nn.Parameter(torch.FloatTensor(1).fill_(10), requires_grad=True) # backbone self.feature = backbone final_feat_dim = 384 self.classifier = nn.Linear(final_feat_dim, 64) self.loss_fn = nn.CrossEntropyLoss() def cos_classifier(self, w, f): """ w.shape = B, nC, d f.shape = B, M, d """ f = F.normalize(f, p=2, dim=f.dim()-1, eps=1e-12) w = F.normalize(w, p=2, dim=w.dim()-1, eps=1e-12) cls_scores = f @ w.transpose(1, 2) # B, M, nC cls_scores = self.scale_cls * (cls_scores + self.bias) return cls_scores def forward(self, supp_x, supp_y, x): """ supp_x.shape = [B, nSupp, C, H, W] supp_y.shape = [B, nSupp] x.shape = [B, nQry, C, H, W] """ num_classes = supp_y.max() + 1 # NOTE: assume B==1 B, nSupp, C, H, W = supp_x.shape supp_f = self.feature.forward(supp_x.contiguous().view(-1, C, H, W)) supp_f = supp_f.view(B, nSupp, -1) supp_y_1hot = F.one_hot(supp_y, num_classes).transpose(1, 2) # B, nC, nSupp # B, nC, nSupp x B, nSupp, d = B, nC, d prototypes = torch.bmm(supp_y_1hot.float(), supp_f) prototypes = prototypes / supp_y_1hot.sum(dim=2, keepdim=True) # NOTE: may div 0 if some classes got 0 images feat = self.feature.forward(x.view(-1, C, H, W)) feat = feat.view(B, x.shape[1], -1) # B, nQry, d logits = self.cos_classifier(prototypes, feat) # B, nQry, nC return logits def set_statues_of_modules(self, flag): if(flag=='eval'): self.feature.eval() self.classifier.eval() #self.scale_cls.eval() #self.bias.eval() elif(flag=='train'): self.feature.train() self.classifier.train() #self.scale_cls.train() #self.bias.train() return def forward_protonet(self, episode_f,supp_y, B, nSupp, nQuery, num_classes): #print('episode_f:', episode_f.size()) episode_f = episode_f.view(num_classes, nSupp + nQuery, -1) #print('episode_f:', episode_f.size()) fea_dim = episode_f.size()[-1] supp_f = episode_f[:, :nSupp, :].contiguous().view(-1, fea_dim).unsqueeze(0) query_f = episode_f[:, nSupp:, :].contiguous().view(-1, fea_dim).unsqueeze(0) supp_y_1hot = F.one_hot(supp_y, num_classes).transpose(1, 2) # B, nC, nSupp # B, nC, nSupp x B, nSupp, d = B, nC, d prototypes = torch.bmm(supp_y_1hot.float(), supp_f) prototypes = prototypes / supp_y_1hot.sum(dim=2, keepdim=True) # NOTE: may div 0 if some classes got 0 images logits = self.cos_classifier(prototypes, query_f) # B, nQry, nC return logits def adversarial_attack_Incre(self, x_ori, y_ori, epsilon_list): x_ori = x_ori.cuda() y_ori = y_ori.cuda() x_size = x_ori.size() x_ori = x_ori.view(x_size[0]*x_size[1], x_size[2], x_size[3], x_size[4]) y_ori = y_ori.view(x_size[0]*x_size[1]) # if not adv, set defalut = 'None' adv_style_mean_block1, adv_style_std_block1 = 'None', 'None' adv_style_mean_block2, adv_style_std_block2 = 'None', 'None' adv_style_mean_block3, adv_style_std_block3 = 'None', 'None' # forward and set the grad = True blocklist = 'block123' if('1' in blocklist and epsilon_list[0] != 0 ): x_ori_block1 = self.feature.forward_block1(x_ori) x_ori_block1_cls, x_ori_block1_P = preprocessing(x_ori_block1) feat_size_block1 = x_ori_block1_P.size() #print('x_ori_block1:', x_ori_block1.size(), x_ori_block1_P.size()) ori_style_mean_block1, ori_style_std_block1 = calc_mean_std(x_ori_block1_P) # set them as learnable parameters ori_style_mean_block1 = torch.nn.Parameter(ori_style_mean_block1) ori_style_std_block1 = torch.nn.Parameter(ori_style_std_block1) ori_style_mean_block1.requires_grad_() ori_style_std_block1.requires_grad_() # contain ori_style_mean_block1 in the graph x_normalized_block1 = (x_ori_block1_P - ori_style_mean_block1.detach().expand(feat_size_block1)) / ori_style_std_block1.detach().expand(feat_size_block1) x_ori_block1_P = x_normalized_block1 * ori_style_std_block1.expand(feat_size_block1) + ori_style_mean_block1.expand(feat_size_block1) x_ori_block1 = postprocessing(x_ori_block1_cls, x_ori_block1_P) #print('x_ori_block1:', x_ori_block1.size()) # pass the rest model x_ori_block2 = self.feature.forward_block2(x_ori_block1) x_ori_block3 = self.feature.forward_block3(x_ori_block2) x_ori_block4 = self.feature.forward_block4(x_ori_block3) x_ori_fea = self.feature.forward_rest(x_ori_block4) x_ori_output = self.classifier.forward(x_ori_fea) # calculate initial pred, loss and acc ori_pred = x_ori_output.max(1, keepdim=True)[1] ori_loss = self.loss_fn(x_ori_output, y_ori) ori_acc = (ori_pred == y_ori).type(torch.float).sum().item() / y_ori.size()[0] # zero all the existing gradients self.feature.zero_grad() self.classifier.zero_grad() # backward loss ori_loss.backward() # collect datagrad grad_ori_style_mean_block1 = ori_style_mean_block1.grad.detach() grad_ori_style_std_block1 = ori_style_std_block1.grad.detach() # fgsm style attack index = torch.randint(0, len(epsilon_list), (1, ))[0] epsilon = epsilon_list[index] adv_style_mean_block1 = fgsm_attack(ori_style_mean_block1, epsilon, grad_ori_style_mean_block1) adv_style_std_block1 = fgsm_attack(ori_style_std_block1, epsilon, grad_ori_style_std_block1) # add zero_grad self.feature.zero_grad() self.classifier.zero_grad() if('2' in blocklist and epsilon_list[1] != 0): x_ori_block1 = self.feature.forward_block1(x_ori) # update adv_block1 x_adv_block1 = changeNewAdvStyle_ViT(x_ori_block1, adv_style_mean_block1, adv_style_std_block1, p_thred=0) # forward block2 x_ori_block2 = self.feature.forward_block2(x_adv_block1) # calculate mean and std x_ori_block2_cls , x_ori_block2_P = preprocessing(x_ori_block2) feat_size_block2 = x_ori_block2_P.size() ori_style_mean_block2, ori_style_std_block2 = calc_mean_std(x_ori_block2_P) # set them as learnable parameters ori_style_mean_block2 = torch.nn.Parameter(ori_style_mean_block2) ori_style_std_block2 = torch.nn.Parameter(ori_style_std_block2) ori_style_mean_block2.requires_grad_() ori_style_std_block2.requires_grad_() # contain ori_style_mean_block1 in the graph x_normalized_block2 = (x_ori_block2_P - ori_style_mean_block2.detach().expand(feat_size_block2)) / ori_style_std_block2.detach().expand(feat_size_block2) x_ori_block2_P = x_normalized_block2 * ori_style_std_block2.expand(feat_size_block2) + ori_style_mean_block2.expand(feat_size_block2) x_ori_block2 = postprocessing(x_ori_block2_cls, x_ori_block2_P) # pass the rest model x_ori_block3 = self.feature.forward_block3(x_ori_block2) x_ori_block4 = self.feature.forward_block4(x_ori_block3) x_ori_fea = self.feature.forward_rest(x_ori_block4) x_ori_output = self.classifier.forward(x_ori_fea) # calculate initial pred, loss and acc ori_pred = x_ori_output.max(1, keepdim=True)[1] ori_loss = self.loss_fn(x_ori_output, y_ori) ori_acc = (ori_pred == y_ori).type(torch.float).sum().item() / y_ori.size()[0] #print('ori_pred:', ori_pred, 'ori_loss:', ori_loss, 'ori_acc:', ori_acc) # zero all the existing gradients self.feature.zero_grad() self.classifier.zero_grad() # backward loss ori_loss.backward() # collect datagrad grad_ori_style_mean_block2 = ori_style_mean_block2.grad.detach() grad_ori_style_std_block2 = ori_style_std_block2.grad.detach() # fgsm style attack index = torch.randint(0, len(epsilon_list), (1, ))[0] epsilon = epsilon_list[index] adv_style_mean_block2 = fgsm_attack(ori_style_mean_block2, epsilon, grad_ori_style_mean_block2) adv_style_std_block2 = fgsm_attack(ori_style_std_block2, epsilon, grad_ori_style_std_block2) #print('adv_style_mean_block2:', adv_style_mean_block2.size(), 'adv_style_std_block2:', adv_style_std_block2.size()) # add zero_grad self.feature.zero_grad() self.classifier.zero_grad() if('3' in blocklist and epsilon_list[2] != 0): x_ori_block1 = self.feature.forward_block1(x_ori) x_adv_block1 = changeNewAdvStyle_ViT(x_ori_block1, adv_style_mean_block1, adv_style_std_block1, p_thred=0) x_ori_block2 = self.feature.forward_block2(x_adv_block1) x_adv_block2 = changeNewAdvStyle_ViT(x_ori_block2, adv_style_mean_block2, adv_style_std_block2, p_thred=0) x_ori_block3 = self.feature.forward_block3(x_adv_block2) x_ori_block3_cls, x_ori_block3_P = preprocessing(x_ori_block3) # calculate mean and std feat_size_block3 = x_ori_block3_P.size() ori_style_mean_block3, ori_style_std_block3 = calc_mean_std(x_ori_block3_P) # set them as learnable parameters ori_style_mean_block3 = torch.nn.Parameter(ori_style_mean_block3) ori_style_std_block3 = torch.nn.Parameter(ori_style_std_block3) ori_style_mean_block3.requires_grad_() ori_style_std_block3.requires_grad_() # contain ori_style_mean_block3 in the graph x_normalized_block3 = (x_ori_block3_P - ori_style_mean_block3.detach().expand(feat_size_block3)) / ori_style_std_block3.detach().expand(feat_size_block3) x_ori_block3_P = x_normalized_block3 * ori_style_std_block3.expand(feat_size_block3) + ori_style_mean_block3.expand(feat_size_block3) x_ori_block3 = postprocessing(x_ori_block3_cls, x_ori_block3_P) # pass the rest model x_ori_block4 = self.feature.forward_block4(x_ori_block3) x_ori_fea = self.feature.forward_rest(x_ori_block4) x_ori_output = self.classifier.forward(x_ori_fea) # calculate initial pred, loss and acc ori_pred = x_ori_output.max(1, keepdim=True)[1] ori_loss = self.loss_fn(x_ori_output, y_ori) ori_acc = (ori_pred == y_ori).type(torch.float).sum().item() / y_ori.size()[0] # zero all the existing gradients self.feature.zero_grad() self.classifier.zero_grad() # backward loss ori_loss.backward() # collect datagrad grad_ori_style_mean_block3 = ori_style_mean_block3.grad.detach() grad_ori_style_std_block3 = ori_style_std_block3.grad.detach() # fgsm style attack index = torch.randint(0, len(epsilon_list), (1, ))[0] epsilon = epsilon_list[index] adv_style_mean_block3 = fgsm_attack(ori_style_mean_block3, epsilon, grad_ori_style_mean_block3) adv_style_std_block3 = fgsm_attack(ori_style_std_block3, epsilon, grad_ori_style_std_block3) return adv_style_mean_block1, adv_style_std_block1, adv_style_mean_block2, adv_style_std_block2, adv_style_mean_block3, adv_style_std_block3 def set_forward_loss_StyAdv(self, SupportTensor,QueryTensor,SupportLabel, QueryLabel, GlobalID_S,GlobalID_Q, epsilon_list): ################################################################## ''' supp_x.shape = [B, nSupp, C, H, W] supp_y.shape = [B, nSupp] x.shape = [B, nQry, C, H, W] # to tacke the input data x_ori: [5, 21, 3, 224, 224], global_y: [5, 21] ''' # to resize as x_ori: torch.Size([5, 21, 3, 224, 224]) global_y: torch.Size([5, 21]) B = SupportTensor.size()[0] num_classes = SupportLabel.max() + 1 # NOTE: assume B==1 SupportTensor = SupportTensor.squeeze().view(num_classes, -1, 3, 224, 224) QueryTensor = QueryTensor.squeeze().view(num_classes, -1, 3, 224, 224) nSupp = SupportTensor.size()[1] nQuery = QueryTensor.size()[1] x_ori = torch.cat((SupportTensor, QueryTensor), dim=1) global_y = torch.cat((GlobalID_S.view(num_classes, nSupp), GlobalID_Q.view(num_classes, nQuery)), dim=1) #print('x_ori:', x_ori.size(), 'global_y:', global_y.size()) ################################################################## # 0. first cp x_adv from x_ori x_adv = x_ori # 1. styleAdv self.set_statues_of_modules('eval') adv_style_mean_block1, adv_style_std_block1, adv_style_mean_block2, adv_style_std_block2, adv_style_mean_block3, adv_style_std_block3 = self.adversarial_attack_Incre(x_ori, global_y, epsilon_list) self.feature.zero_grad() self.classifier.zero_grad() # 2. forward and get loss self.set_statues_of_modules('train') x_ori = x_ori.cuda() x_size = x_ori.size() x_ori = x_ori.view(num_classes*(nSupp+nQuery), 3, 224, 224) global_y = global_y.view(num_classes*(nSupp+nQuery)).cuda() x_ori_block1 = self.feature.forward_block1(x_ori) x_ori_block2 = self.feature.forward_block2(x_ori_block1) x_ori_block3 = self.feature.forward_block3(x_ori_block2) x_ori_block4 = self.feature.forward_block4(x_ori_block3) x_ori_fea = self.feature.forward_rest(x_ori_block4) # 3. ori cls global loss scores_cls_ori = self.classifier.forward(x_ori_fea) loss_cls_ori = self.loss_fn(scores_cls_ori, global_y) # 4. ori FSL scores and losses scores_fsl_ori = self.forward_protonet(x_ori_fea, SupportLabel,B, nSupp, nQuery, num_classes) scores_fsl_ori = scores_fsl_ori.view(num_classes*nQuery, -1) QueryLabel = QueryLabel.view(-1) loss_fsl_ori = self.loss_fn(scores_fsl_ori, QueryLabel) # 5. forward StyleAdv x_adv = x_adv.cuda() x_adv = x_adv.view(x_size[0]*x_size[1], x_size[2], x_size[3], x_size[4]) x_adv_block1 = self.feature.forward_block1(x_adv) x_adv_block1_newStyle = changeNewAdvStyle_ViT(x_adv_block1, adv_style_mean_block1, adv_style_std_block1, p_thred = P_THRED) x_adv_block2 = self.feature.forward_block2(x_adv_block1_newStyle) x_adv_block2_newStyle = changeNewAdvStyle_ViT(x_adv_block2, adv_style_mean_block2, adv_style_std_block2, p_thred = P_THRED) x_adv_block3 = self.feature.forward_block3(x_adv_block2_newStyle) x_adv_block3_newStyle = changeNewAdvStyle_ViT(x_adv_block3, adv_style_mean_block3, adv_style_std_block3, p_thred = P_THRED) x_adv_block4 = self.feature.forward_block4(x_adv_block3_newStyle) x_adv_fea = self.feature.forward_rest(x_adv_block4) # 6. adv cls gloabl loss scores_cls_adv = self.classifier.forward(x_adv_fea) loss_cls_adv = self.loss_fn(scores_cls_adv, global_y) # 7. adv FSL scores and losses scores_fsl_adv = self.forward_protonet(x_adv_fea, SupportLabel,B, nSupp, nQuery, num_classes) scores_fsl_adv = scores_fsl_adv.view(num_classes*nQuery, -1) loss_fsl_adv = self.loss_fn(scores_fsl_adv, QueryLabel) return scores_fsl_ori, loss_fsl_ori, scores_cls_ori, loss_cls_ori, scores_fsl_adv, loss_fsl_adv, scores_cls_adv, loss_cls_adv