|
|
import torch |
|
|
from torch import nn |
|
|
import numpy as np |
|
|
|
|
|
def l2_norm(input,axis=1): |
|
|
norm = torch.norm(input, 2, axis, True) |
|
|
output = torch.div(input, norm) |
|
|
return output |
|
|
|
|
|
class L2_normalization(nn.Module): |
|
|
def forward(self, input): |
|
|
return l2_norm(input) |
|
|
|
|
|
def freeze_bert_parameters(model): |
|
|
for name, param in model.bert.named_parameters(): |
|
|
param.requires_grad = False |
|
|
if "encoder.layer.11" in name or "pooler" in name: |
|
|
param.requires_grad = True |
|
|
return model |
|
|
|
|
|
def freeze_bert_parameters_KCL(model): |
|
|
for name, param in model.encoder_q.named_parameters(): |
|
|
param.requires_grad = False |
|
|
if "encoder.layer.11" in name or "pooler" in name: |
|
|
param.requires_grad = True |
|
|
for name, param in model.encoder_k.named_parameters(): |
|
|
param.requires_grad = False |
|
|
if "encoder.layer.11" in name or "pooler" in name: |
|
|
param.requires_grad = True |
|
|
return model |
|
|
|
|
|
class ConvexSampler(nn.Module): |
|
|
def __init__(self, args): |
|
|
super(ConvexSampler, self).__init__() |
|
|
self.multiple_convex = args.multiple_convex |
|
|
self.multiple_convex_eval = args.multiple_convex_eval |
|
|
self.unseen_label_id = args.unseen_label_id |
|
|
self.device = args.device |
|
|
self.batch_size = args.train_batch_size |
|
|
self.oos_num = args.train_batch_size |
|
|
self.feat_dim = args.feat_dim |
|
|
|
|
|
def forward(self, z, label_ids, mode=None): |
|
|
num_convex = self.batch_size * self.multiple_convex |
|
|
num_convex_eval = self.batch_size * self.multiple_convex_eval |
|
|
convex_list = [] |
|
|
if mode =='train': |
|
|
if label_ids.size(0)>2: |
|
|
while len(convex_list) < num_convex: |
|
|
cdt = np.random.choice(label_ids.size(0), 2, replace=False) |
|
|
|
|
|
if label_ids[cdt[0]] != label_ids[cdt[1]]: |
|
|
s = np.random.uniform(0, 1, 1) |
|
|
convex_list.append(s[0] * z[cdt[0]] + (1 - s[0]) * z[cdt[1]]) |
|
|
convex_samples = torch.cat(convex_list, dim=0).view(num_convex, -1) |
|
|
z = torch.cat((z, convex_samples), dim=0) |
|
|
label_ids = torch.cat((label_ids, torch.tensor([self.unseen_label_id] * num_convex).to(self.device)), dim=0) |
|
|
elif mode == 'eval': |
|
|
if label_ids.size(0) > 2: |
|
|
val_num = num_convex_eval |
|
|
while len(convex_list) < val_num: |
|
|
cdt = np.random.choice(label_ids.size(0), 2, replace=False) |
|
|
if label_ids[cdt[0]] != label_ids[cdt[1]]: |
|
|
s = np.random.uniform(0, 1, 1) |
|
|
convex_list.append(s[0] * z[cdt[0]] + (1 - s[0]) * z[cdt[1]]) |
|
|
convex_samples = torch.cat(convex_list, dim=0).view(val_num, -1) |
|
|
z = torch.cat((z, convex_samples), dim=0) |
|
|
label_ids = torch.cat((label_ids, torch.tensor([self.unseen_label_id] * val_num).to(self.device)), dim=0) |
|
|
return z, label_ids |
|
|
|
|
|
def pair_cosine_similarity(x, x_adv, eps=1e-8): |
|
|
n = x.norm(p=2, dim=1, keepdim=True) |
|
|
n_adv = x_adv.norm(p=2, dim=1, keepdim=True) |
|
|
return (x @ x.t()) / (n * n.t()).clamp(min=eps), (x_adv @ x_adv.t()) / (n_adv * n_adv.t()).clamp(min=eps), (x @ x_adv.t()) / (n * n_adv.t()).clamp(min=eps) |
|
|
|
|
|
def nt_xent(x, x_adv, mask, cuda=True, t=0.1): |
|
|
x, x_adv, x_c = pair_cosine_similarity(x, x_adv) |
|
|
x = torch.exp(x / t) |
|
|
x_adv = torch.exp(x_adv / t) |
|
|
x_c = torch.exp(x_c / t) |
|
|
mask_count = mask.sum(1) |
|
|
mask_reverse = (~(mask.bool())).long() |
|
|
if cuda: |
|
|
dis = (x * (mask - torch.eye(x.size(0)).long().cuda()) + x_c * mask) / (x.sum(1) + x_c.sum(1) - torch.exp(torch.tensor(1 / t))) + mask_reverse |
|
|
dis_adv = (x_adv * (mask - torch.eye(x.size(0)).long().cuda()) + x_c.T * mask) / (x_adv.sum(1) + x_c.sum(0) - torch.exp(torch.tensor(1 / t))) + mask_reverse |
|
|
else: |
|
|
dis = (x * (mask - torch.eye(x.size(0)).long()) + x_c * mask) / (x.sum(1) + x_c.sum(1) - torch.exp(torch.tensor(1 / t))) + mask_reverse |
|
|
dis_adv = (x_adv * (mask - torch.eye(x.size(0)).long()) + x_c.T * mask) / (x_adv.sum(1) + x_c.sum(0) - torch.exp(torch.tensor(1 / t))) + mask_reverse |
|
|
loss = (torch.log(dis).sum(1) + torch.log(dis_adv).sum(1)) / mask_count |
|
|
return -loss.mean() |
|
|
|
