def _train_a_query(self, indice): indexer = np.stack( [np.array(indice[i::self.num_groups]) for i in range(self.num_groups)]) sampler = SubsetRandomSampler(indice) loader = DataLoader(self.dataset, batch_size=128, sampler=sampler, num_workers=4, pin_memory=True) alpha_generator = Exponential(torch.ones([1, self.num_groups])) t_iter = tqdm(range(self.num_epoch), total=self.num_epoch, desc="Training") for t in t_iter: alpha = alpha_generator.sample().cuda() for img, label, index in loader: self.model.train() n0 = img.size(0) u_is = [] for i in index: u_i = np.where(indexer == i.item())[0][0] u_is += [u_i] w = alpha[0, u_is].cuda() output = self.model(img.cuda(), alpha.repeat_interleave(n0, 0)) loss = self.loss_fn(output, label.cuda(), w) self.optim.zero_grad() loss.backward() self.optim.step() self.sched.step()
def odin_infer(loader, model, num_bs, num_classes, with_acc=False, seed=0, T=1000, eps=0.0014): loss_fn = torch.nn.CrossEntropyLoss() torch.manual_seed(seed) model.eval() a_test_ = Exponential(torch.ones([1, 400])) a_test = a_test_.sample((num_bs, )) acc = 0. outputs = np.zeros([num_bs, len(loader.dataset), num_classes + 1]) beg = 0 for i, (img, label) in enumerate(loader): index = list(range(beg, beg + img.size(0))) beg = beg + img.size(0) label = label.numpy().reshape(-1, 1) img_ = img.cuda() img_.requires_grad = True output = model(img_, torch.zeros([img.shape[0], 400])) output = output / T pseudo_label = output.argmax(-1).cuda() loss = loss_fn(output, pseudo_label) loss.backward() gradient = torch.ge(img_.grad.data, 0) gradient = (gradient.float() - 0.5) * 2 gradient.index_copy_( 1, torch.LongTensor([0]).cuda(), gradient.index_select(1, torch.LongTensor([0]).cuda()) / (0.2023)) gradient.index_copy_( 1, torch.LongTensor([1]).cuda(), gradient.index_select(1, torch.LongTensor([1]).cuda()) / (0.1994)) gradient.index_copy_( 1, torch.LongTensor([2]).cuda(), gradient.index_select(1, torch.LongTensor([2]).cuda()) / (0.2010)) img_new = torch.add(img_.data, -eps, gradient) for _ in range(num_bs): w_test = a_test[_].repeat_interleave(img.shape[0], dim=0) output_new = model(img_new, w_test).cpu().detach().numpy() outputs[_, index] = np.concatenate([output_new, label], axis=1) if with_acc: pred = outputs.sum(0)[:, :-1].argmax(1) label = outputs[0][:, -1] acc = pred == label print(f"[Test] acc : {acc.mean()}") return outputs
def infer_a_sample(img, model, num_classes, num_bs, fac): model.eval() a_test = Exponential(torch.ones([1, 400])) w_test = a_test.sample((num_bs, )) output = np.zeros([num_bs, num_classes]) for _ in range(num_bs): o = model(img, w_test[_], fac).cpu().numpy() output[_] = o return output
class HomogeneousPoissonProcess: def __init__(self, rate=1): self.rate = rate self.exp = Exponential(rate) def sample(self, size, max_seq_len, max_time=math.inf): gaps = self.exp.sample((size, max_seq_len)) times = torch.cumsum(gaps, dim=1) masks = (times <= max_time).float() return times, masks
def __init__(self, in_features, alpha=None, alpha_trainable=True): """ Init method. """ super(Snake, self).__init__() self.in_features = in_features if alpha is not None: self.alpha = Parameter(torch.ones(in_features) * alpha) else: m = Exponential(torch.tensor([0.1])) self.alpha = Parameter(m.sample(in_features)) self.alpha.requiresGrad = alpha_trainable
def _infer_gbs(self, with_acc=False, seed=0): torch.manual_seed(seed) a_test_ = Exponential(torch.ones([1, self.args.n_a])) a_test = a_test_.sample([self.args.num_bs, ]) outputs = np.zeros([self.args.num_bs, len(self.loader.dataset), self.args.num_classes + 1]) beg = 0 for i, (img, label) in enumerate(self.loader): end = beg + img.size(0) label = label.numpy().reshape(-1, 1) for _ in range(self.args.num_bs): w_test = a_test[_].repeat_interleave(img.size(0), dim=0) output = self._infer_a_batch_a_bs(img, w_test) outputs[_, beg: end] = np.concatenate([output, label], axis=1) beg = end if with_acc: pred = outputs.sum(0)[:, :-1].argmax(1) label = outputs[0][:, -1] acc = (pred == label).mean() print(f"[Test] acc : {acc}") return outputs
def odin_infer_a_sample(img, model, num_classes, num_bs, fac, T=1000, eps=0.0001): model.eval() loss_fn = torch.nn.CrossEntropyLoss() img = img.cuda() img.requires_grad = True model.zero_grad() a_test = Exponential(torch.ones([1, 400])) w_test = a_test.sample((num_bs, )) outputs = np.zeros([num_bs, num_classes]) for _ in range(num_bs): # w_test = a_test[_].repeat_interleave(img.shape[0], dim=0) img_ = img.cuda() img_.requires_grad = True output = model(img_, w_test[_], fac) output = output / T pseudo_label = output.argmax(-1).cuda() loss = loss_fn(output, pseudo_label) loss.backward() gradient = torch.ge(img_.grad.data, 0) # gradient = (gradient.float() - 0.5) * 2 # gradient.index_copy_(1, torch.LongTensor([0]).cuda(), # gradient.index_select(1, torch.LongTensor([0]).cuda()) / (0.2023)) # gradient.index_copy_(1, torch.LongTensor([1]).cuda(), # gradient.index_select(1, torch.LongTensor([1]).cuda()) / (0.1994)) # gradient.index_copy_(1, torch.LongTensor([2]).cuda(), # gradient.index_select(1, torch.LongTensor([2]).cuda()) / (0.2010)) img_new = torch.add(img_.data, -eps, gradient) output_new = model(img_new, w_test[_], fac).cpu().detach().numpy() outputs[_] = output_new return outputs
def infer(loader, model, num_bs, num_classes, with_acc=False, with_indice=False, is_mcd=False, seed=0): torch.manual_seed(seed) model.eval() a_test_ = Exponential(torch.ones([1, 400])) a_test = a_test_.sample((num_bs, )) acc = 0. outputs = np.zeros([num_bs, len(loader.dataset), num_classes + 1]) beg = 0 indice = [] ret = [] for i, (img, label, _) in enumerate(loader): img = img.cuda() index = list(range(beg, beg + img.size(0))) beg = beg + img.size(0) label = label.numpy().reshape(-1, 1) indice += [_] for _ in range(num_bs): w_test = a_test[_].repeat_interleave(img.shape[0], dim=0) if is_mcd: model.apply(apply_dropout) output = model(img, w_test.cuda()).cpu().numpy() outputs[_, index] = np.concatenate([output, label], axis=1) ret += [outputs] if with_acc: pred = outputs.sum(0)[:, :-1].argmax(1) label = outputs[0][:, -1] acc = pred == label print(f"[Test] acc : {acc.mean()}") ret += [acc.mean()] if with_indice: ret += [torch.cat(indice)] return ret
def __init__(self, in_features, alpha = None, alpha_trainable = True): ''' Initialization. INPUT: - in_features: shape of the input - alpha: trainable parameter alpha is initialized to 1 by default, higher values = higher-frequency, 5-50 is a good starting point if you already think your data is periodic, consider starting lower e.g. 0.5 if you think not, but don't worry, alpha will be trained along with the rest of your model. ''' super(Snake,self).__init__() self.in_features = in_features # initialize alpha if alpha is not None: self.alpha = Parameter(torch.ones(in_features) * alpha) # create a tensor out of alpha else: m = Exponential(torch.tensor([0.1])) self.alpha = Parameter(m.sample(in_features)) # Random init = mix of frequencies self.alpha.requiresGrad = alpha_trainable # Usually we'll want to train alpha, but maybe for some experiments we won't?