Python torch.argmax() Examples

def test_softmax(self):
        em = LogisticRegression(seed=1, input_dim=2, output_dim=3, verbose=False)
        Xs, _ = self.single_problem
        Ys = []
        for X in Xs:
            class1 = X[:, 0] < X[:, 1]
            class2 = X[:, 0] > X[:, 1] + 0.5
            class3 = X[:, 0] > X[:, 1]
            Y = torch.argmax(torch.stack([class1, class2, class3], dim=1), dim=1) + 1
            Ys.append(Y)
        em.train_model(
            (Xs[0], Ys[0]),
            valid_data=(Xs[1], Ys[1]),
            lr=0.1,
            n_epochs=10,
            checkpoint=False,
        )
        score = em.score((Xs[2], Ys[2]), verbose=False)
        self.assertGreater(score, 0.95) 

def train_cnn(net, train_iter, test_iter, batch_size, optimizer, device, num_epochs):
    net = net.to(device)
    print('training on', device)
    loss = nn.CrossEntropyLoss()
    batch_count = 0
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n, start = 0.0, 0.0, 0, time.time()
        for X, y in train_iter:
            X = X.to(device)
            y = y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            optimizer.zero_grad()
            l.backward()
            optimizer.step()

            train_l_sum += l.cpu().item()
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().cpu().item()
            n += y.shape[0]
            batch_count += 1
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f, time %.1f sec' %
              (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc, time.time() - start)) 

def _sample_action(self, cat_distr, mask, relaxed, tau_weights, straight_through, gumbel_noise):
        if self.training:
            if relaxed:
                N = mask.sum(dim=-1, keepdim=True)
                tau = tau_weights[0] + tau_weights[1].exp() * torch.log(N + 1) + tau_weights[2].exp() * N
                actions, gumbel_noise = cat_distr.rsample(temperature=tau, gumbel_noise=gumbel_noise)
                if straight_through:
                    actions_hard = torch.zeros_like(actions)
                    actions_hard.scatter_(-1, actions.argmax(dim=-1, keepdim=True), 1.0)
                    actions = (actions_hard - actions).detach() + actions
                actions = clamp_grad(actions, -0.5, 0.5)
            else:
                actions, gumbel_noise = cat_distr.rsample(gumbel_noise=gumbel_noise)
        else:
            actions = torch.zeros_like(cat_distr.probs)
            actions.scatter_(-1, torch.argmax(cat_distr.probs, dim=-1, keepdim=True), 1.0)
            gumbel_noise = None
        return actions, gumbel_noise 

def batch_intersection_union(output, target, nclass):
    """mIoU"""
    # inputs are numpy array, output 4D, target 3D
    mini = 1
    maxi = nclass
    nbins = nclass
    predict = torch.argmax(output, 1) + 1
    target = target.float() + 1

    predict = predict.float() * (target > 0).float()
    intersection = predict * (predict == target).float()
    # areas of intersection and union
    # element 0 in intersection occur the main difference from np.bincount. set boundary to -1 is necessary.
    area_inter = torch.histc(intersection.cpu(), bins=nbins, min=mini, max=maxi)
    area_pred = torch.histc(predict.cpu(), bins=nbins, min=mini, max=maxi)
    area_lab = torch.histc(target.cpu(), bins=nbins, min=mini, max=maxi)
    area_union = area_pred + area_lab - area_inter
    assert torch.sum(area_inter > area_union).item() == 0, "Intersection area should be smaller than Union area"
    return area_inter.float(), area_union.float() 

def BPDA_attack(image,target, model, step_size = 1., iterations = 10, linf=False, transform_func=identity_transform):
    target = label2tensor(target)
    adv = image.detach().numpy()
    adv = torch.from_numpy(adv)
    adv.requires_grad_()
    for _ in range(iterations):
        adv_def = transform_func(adv)
        adv_def.requires_grad_()
        l2 = nn.MSELoss()
        loss = l2(0, adv_def)
        loss.backward()
        g = get_cw_grad(adv_def, image, target, model)
        if linf:
            g = torch.sign(g)
        print(g.numpy().sum())
        adv = adv.detach().numpy() - step_size * g.numpy()
        adv = clip_bound(adv)
        adv = torch.from_numpy(adv)
        adv.requires_grad_()
        if linf:
            print('label', torch.argmax(model(adv)), 'linf', torch.max(torch.abs(adv - image)).detach().numpy())
        else:
            print('label', torch.argmax(model(adv)), 'l2', l2_norm(adv, image))
    return adv.detach().numpy() 

def _graph_fn_get_action_components(self, flat_key, logits, parameters, deterministic):
        action_space_component = self.flat_action_space[flat_key]

        # Skip our distribution, iff discrete action-space and deterministic acting (greedy).
        # In that case, one does not need to create a distribution in the graph each act (only to get the argmax
        # over the logits, which is the same as the argmax over the probabilities (or log-probabilities)).
        if isinstance(action_space_component, IntBox) and \
                (deterministic is True or (isinstance(deterministic, np.ndarray) and deterministic)):
            return self._graph_fn_get_deterministic_action_wo_distribution(logits)
        # Bernoulli: Sigmoid derived p must be larger 0.5.
        elif isinstance(action_space_component, BoolBox) and \
                (deterministic is True or (isinstance(deterministic, np.ndarray) and deterministic)):
            # Note: Change 0.5 to 1.0, once parameters are logits, not probs anymore (so far, parameters for
            # Bernoulli distributions are still probs).
            if get_backend() == "tf":
                return tf.greater(parameters, 0.5)
            elif get_backend() == "pytorch":
                return torch.gt(parameters, 0.5)
        # Deterministic is tensor or False. Pass through graph.
        else:
            return self.distributions[flat_key].draw(parameters, deterministic) 

def batch_intersection_union(output, target, nclass):
    """mIoU"""
    # inputs are NDarray, output 4D, target 3D
    # the category -1 is ignored class, typically for background / boundary
    mini = 1
    maxi = nclass
    nbins = nclass
    predict = torch.argmax(output, 1) + 1
    target = target.float() + 1

    predict = predict.float() * (target > 0).float()
    intersection = predict * (predict == target).float()
    # areas of intersection and union
    area_inter = torch.histc(intersection, bins=nbins, min=mini, max=maxi)
    area_pred = torch.histc(predict, bins=nbins, min=mini, max=maxi)
    area_lab = torch.histc(target, bins=nbins, min=mini, max=maxi)
    area_union = area_pred + area_lab - area_inter
    assert torch.sum(area_inter > area_union).item() == 0, \
        "Intersection area should be smaller than Union area"
    return area_inter.float(), area_union.float() 

def classwise_f1(output, gt):
    """
    Args:
        output: torch.Tensor of shape (n_batch, n_classes, image.shape)
        gt: torch.LongTensor of shape (n_batch, image.shape)
    """

    epsilon = 1e-20
    n_classes = output.shape[1]

    output = torch.argmax(output, dim=1)
    true_positives = torch.tensor([((output == i) * (gt == i)).sum() for i in range(n_classes)]).float()
    selected = torch.tensor([(output == i).sum() for i in range(n_classes)]).float()
    relevant = torch.tensor([(gt == i).sum() for i in range(n_classes)]).float()

    precision = (true_positives + epsilon) / (selected + epsilon)
    recall = (true_positives + epsilon) / (relevant + epsilon)
    classwise_f1 = 2 * (precision * recall) / (precision + recall)

    return classwise_f1 

def _graph_fn_sample_deterministic(self, distribution):
        """
        Returns the argmax (int) of a relaxed one-hot vector. See `_graph_fn_sample_stochastic` for details.
        """
        if get_backend() == "tf":
            # Cast to float again because this is called from a tf.cond where the other option calls a stochastic
            # sample returning a float.
            argmax = tf.argmax(input=distribution._distribution.probs, axis=-1, output_type=tf.int32)
            sample = tf.cast(argmax, dtype=tf.float32)
            # Argmax turns (?, n) into (?,), not (?, 1)
            # TODO: What if we have a time rank as well?
            if len(sample.shape) == 1:
                sample = tf.expand_dims(sample, -1)
            return sample
        elif get_backend() == "pytorch":
            # TODO: keepdims?
            return torch.argmax(distribution.probs, dim=-1).int() 

def evaluate(dataloader, model, prog_args, logger=None):
    '''
    evaluate function
    '''
    if logger is not None and prog_args.save_dir is not None:
        model.load_state_dict(torch.load(prog_args.save_dir + "/" + prog_args.dataset
                                         + "/model.iter-" + str(logger['best_epoch'])))
    model.eval()
    correct_label = 0
    with torch.no_grad():
        for batch_idx, (batch_graph, graph_labels) in enumerate(dataloader):
            if torch.cuda.is_available():
                for (key, value) in batch_graph.ndata.items():
                    batch_graph.ndata[key] = value.cuda()
                graph_labels = graph_labels.cuda()
            ypred = model(batch_graph)
            indi = torch.argmax(ypred, dim=1)
            correct = torch.sum(indi == graph_labels)
            correct_label += correct.item()
    result = correct_label / (len(dataloader) * prog_args.batch_size)
    return result 

def evaluate(epoch, args, model, feats, labels, train, val, test):
    with torch.no_grad():
        batch_size = args.eval_batch_size
        if batch_size <= 0:
            pred = model(feats)
        else:
            pred = []
            num_nodes = labels.shape[0]
            n_batch = (num_nodes + batch_size - 1) // batch_size
            for i in range(n_batch):
                batch_start = i * batch_size
                batch_end = min((i + 1) * batch_size, num_nodes)
                batch_feats = [feat[batch_start: batch_end] for feat in feats]
                pred.append(model(batch_feats))
            pred = torch.cat(pred)

        pred = torch.argmax(pred, dim=1)
        correct = (pred == labels).float()
        train_acc = correct[train].sum() / len(train)
        val_acc = correct[val].sum() / len(val)
        test_acc = correct[test].sum() / len(test)
        return train_acc, val_acc, test_acc 

def calculate_outputs_and_gradients(inputs, model, target_label_idx, cuda=False):
    # do the pre-processing
    predict_idx = None
    gradients = []
    for input in inputs:
        input = pre_processing(input, cuda)
        output = model(input)
        output = F.softmax(output, dim=1)
        if target_label_idx is None:
            target_label_idx = torch.argmax(output, 1).item()
        index = np.ones((output.size()[0], 1)) * target_label_idx
        index = torch.tensor(index, dtype=torch.int64)
        if cuda:
            index = index.cuda()
        output = output.gather(1, index)
        # clear grad
        model.zero_grad()
        output.backward()
        gradient = input.grad.detach().cpu().numpy()[0]
        gradients.append(gradient)
    gradients = np.array(gradients)
    return gradients, target_label_idx 

def evaluate_accuracy(data_iter, net,
                      device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')):
    acc_sum, n = 0.0, 0
    with torch.no_grad():
        for X, y in data_iter:
            if isinstance(net, torch.nn.Module):
                net.eval() # 评估模式，会关闭 dropout
                acc_sum += (net(X.to(device)).argmax(dim=1) == y.to(device)).float().sum().cpu().item()
                net.train() # 改回训练模式
            else:
                # 如果是自定义的模型
                if 'is_training' in net.__code__.co_varnames:
                    acc_sum += (net(X, is_training=False).argmax(dim=1) == y).float().sum().item()
                else:
                    acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
            n += y.shape[0]
    return acc_sum / n 

def predict_rnn_pytorch(prefix, num_chars, model, vocab_size, device, idx_to_char,
                        char_to_idx):
    state = None
    output = [char_to_idx[prefix[0]]]  # output会记录prefix加上输出
    for t in range(num_chars + len(prefix) - 1):
        X = torch.tensor([output[-1]], device=device).view(1, 1)
        if state is not None:
            if isinstance(state, tuple):  # LSTM, state:(h, c)
                state = (state[0].to(device), state[1].to(device))
            else:
                state = state.to(device)

        (Y, state) = model(X, state)  # 前向计算不需要传入模型参数
        if t < len(prefix) - 1:
            output.append(char_to_idx[prefix[t + 1]])
        else:
            output.append(int(Y.argmax(dim=1).item()))
    return ''.join([idx_to_char[i] for i in output]) 

def val_gzsl(self, test_X, test_label, target_classes):

        with torch.no_grad():
            start = 0
            ntest = test_X.size()[0]
            predicted_label = torch.LongTensor(test_label.size())
            for i in range(0, ntest, self.batch_size):

                end = min(ntest, start+self.batch_size)

                output = self.model(test_X[start:end]) #.to(self.device)

                #_, predicted_label[start:end] = torch.max(output.data, 1)
                predicted_label[start:end] = torch.argmax(output.data, 1)

                start = end

            #print(str(predicted_label[:3]).ljust(40,'.'), end= ' '     )
            acc = self.compute_per_class_acc_gzsl(test_label, predicted_label, target_classes)
            return acc 

def choose_action(self, observation):
        if np.random.random() > self.epsilon:
            state = T.tensor([observation],dtype=T.float).to(self.q_eval.device)
            _, advantage = self.q_eval.forward(state)
            action = T.argmax(advantage).item()
        else:
            action = np.random.choice(self.action_space)

        return action 

def compute_acc(pred, labels):
    """
    Compute the accuracy of prediction given the labels.
    """
    return (th.argmax(pred, dim=1) == labels).float().sum() / len(pred) 

def choose_action(self, observation):
        if np.random.random() > self.epsilon:
            state = T.tensor([observation],dtype=T.float).to(self.q_eval.device)
            actions = self.q_eval.forward(state)
            action = T.argmax(actions).item()
        else:
            action = np.random.choice(self.action_space)

        return action 

def learn(self):
        if self.memory.mem_cntr < self.batch_size:
            return

        self.q_eval.optimizer.zero_grad()

        self.replace_target_network()

        states, actions, rewards, states_, dones = self.sample_memory()
        indices = np.arange(self.batch_size)

        q_pred = self.q_eval.forward(states)[indices, actions]
        q_next = self.q_next.forward(states_)
        q_eval = self.q_eval.forward(states_)

        max_actions = T.argmax(q_eval, dim=1)
        q_next[dones] = 0.0

        q_target = rewards + self.gamma*q_next[indices, max_actions]
        loss = self.q_eval.loss(q_target, q_pred).to(self.q_eval.device)
        loss.backward()

        self.q_eval.optimizer.step()
        self.learn_step_counter += 1

        self.decrement_epsilon() 

def choose_action(self, observation):
        if np.random.random() > self.epsilon:
            state = T.tensor([observation],dtype=T.float).to(self.q_eval.device)
            _, advantage = self.q_eval.forward(state)
            action = T.argmax(advantage).item()
        else:
            action = np.random.choice(self.action_space)

        return action 

def compute_acc(pred, labels):
    """
    Compute the accuracy of prediction given the labels.
    """
    return (th.argmax(pred, dim=1) == labels).float().sum() / len(pred) 

def construct_igraph(self, type_scores, edge_scores, stochastic=True):
        # construct igraphs from node type and edge scores
        # note that when stochastic=True, type_scores should be raw scores before softmax, 
        # and edge_scores should probabilities between [0, 1] (after sigmoid)
        assert(type_scores.shape[:2] == edge_scores.shape[:2])
        if stochastic:
            type_probs = F.softmax(type_scores, 2).cpu().detach().numpy()
        G = []
        for gi in range(len(type_scores)):
            # add vertices
            g = igraph.Graph(directed=True)
            for vj in range(self.max_n):
                if stochastic:
                    new_type = np.random.choice(range(self.nvt), p=type_probs[gi][vj])
                else:
                    new_type = torch.argmax(type_scores[gi][vj], 0).item()
                g.add_vertex(type=new_type)
            # add edges
            output_vertex = None
            for vj in range(self.max_n):
                for ek in range(self.max_n):
                    ek_score = edge_scores[gi][vj][ek].item()
                    if stochastic:
                        if np.random.random_sample() < ek_score:
                            g.add_edge(ek, vj)
                    else:
                        if ek_score > 0.5:
                            g.add_edge(ek, vj)
            # apply a topological order to g (otherwise some validity test doesn't work)
            if g.is_dag():
                topo_order = g.topological_sorting()
                perm = [-1] * len(topo_order)
                for i in range(len(perm)):
                    perm[topo_order[i]] = i
                g = g.permute_vertices(perm) 
            G.append(g)
        return G 

def construct_igraph(self, type_scores, edge_scores, stochastic=True):
        # construct igraphs from node type and edge scores
        # note that when stochastic=True, type_scores should be raw scores before softmax, 
        # and edge_scores should probabilities between [0, 1] (after sigmoid)
        assert(type_scores.shape[:2] == edge_scores.shape[:2])
        if stochastic:
            type_probs = F.softmax(type_scores, 2).cpu().detach().numpy()
        G = []
        for gi in range(len(type_scores)):
            g = igraph.Graph(directed=True)
            g.add_vertex(type=self.START_TYPE)
            for vj in range(1, self.max_n):
                if vj == self.max_n - 1:
                    new_type = self.END_TYPE
                else:
                    if stochastic:
                        new_type = np.random.choice(range(self.nvt), p=type_probs[gi][vj-1])
                    else:
                        new_type = torch.argmax(type_scores[gi][vj-1], 0).item()
                g.add_vertex(type=new_type)
                if new_type == self.END_TYPE:  
                    end_vertices = set([v.index for v in g.vs.select(_outdegree_eq=0) 
                                        if v.index != g.vcount()-1])
                    for v in end_vertices:
                        g.add_edge(v, vj)
                    break
                else:
                    for ek in range(vj):
                        ek_score = edge_scores[gi][vj-1][ek].item()
                        if stochastic:
                            if np.random.random_sample() < ek_score:
                                g.add_edge(ek, vj)
                        else:
                            if ek_score > 0.5:
                                g.add_edge(ek, vj)
            G.append(g)
        return G 

def next(self, Xcandidates):

        if not self.N:
            return torch.randperm(Xcandidates.size()[0])[0]

        fmean, fvar = self.posterior(Xcandidates)

        return torch.argmax(self.acq_func(fmean, fvar, self.ybest)) 

def _graph_fn_get_deterministic_action_wo_distribution(self, logits):
        """
        Use this function only for discrete action spaces to circumvent using a full-blown
        backend-specific distribution object (e.g. tf.distribution.Multinomial).

        Args:
            logits (SingleDataOp): Logits over which to pick the argmax (greedy action).

        Returns:
            SingleDataOp: The argmax over the last rank of the input logits.
        """
        if get_backend() == "tf":
            return tf.argmax(logits, axis=-1, output_type=tf.int32)
        elif get_backend() == "pytorch":
            return torch.argmax(logits, dim=-1).int() 

def forward(self, graph, labels=None):
        etypes = graph.edata.pop('type')
        annotation = graph.ndata.pop('annotation').float()

        assert annotation.size()[-1] == self.annotation_size

        node_num = graph.number_of_nodes()

        zero_pad = torch.zeros([node_num, self.out_feats - self.annotation_size],
                               dtype=torch.float,
                               device=annotation.device)

        h1 = torch.cat([annotation, zero_pad], -1)
        out = self.ggnn(graph, h1, etypes)

        all_logits = self.output_layer(torch.cat([out, annotation], -1)).squeeze(-1)
        graph.ndata['logits'] = all_logits

        batch_g = dgl.unbatch(graph)

        preds = []
        if labels is not None:
            loss = 0.0
        for i, g in enumerate(batch_g):
            logits = g.ndata['logits']
            preds.append(torch.argmax(logits))
            if labels is not None:
                logits = logits.unsqueeze(0)
                y = labels[i].unsqueeze(0)
                loss += self.loss_fn(logits, y)

        if labels is not None:
            loss /= float(len(batch_g))
            return loss, preds
        return preds 

def compute_acc(pred, labels):
    """
    Compute the accuracy of prediction given the labels.
    """
    return (th.argmax(pred, dim=1) == labels).float().sum() / len(pred) 

def _graph_fn_sample_deterministic(self, distribution):
        if get_backend() == "tf":
            return tf.argmax(input=distribution.probs, axis=-1, output_type=util.convert_dtype("int"))
        elif get_backend() == "pytorch":
            return torch.argmax(distribution.probs, dim=-1).int() 

def _graph_fn_sample_stochastic(self, distribution):
        """
        Returns a relaxed one-hot vector representing a quasi-one-hot action.
        To get the actual int action, one would have to take the argmax over
        these output values. However, argmax would break the differentiability and should
        thus only be used right before applying the action in e.g. an env.
        """
        if get_backend() == "tf":
            return distribution.sample(seed=self.seed)
        elif get_backend() == "pytorch":
            return distribution.sample() 

def compute_acc(pred, labels):
    """
    Compute the accuracy of prediction given the labels.
    """
    return (th.argmax(pred, dim=1) == labels).float().sum() / len(pred)