Python chainer.cuda.get_array_module() Examples

def check_forward(self, x_data):
        xp = cuda.get_array_module(x_data)
        y = maximum_entropy_mellowmax(x_data)
        self.assertEqual(y.array.dtype, self.dtype)

        print('y', y.array)

        # Outputs must be positive
        xp.testing.assert_array_less(xp.zeros_like(y.array), y.array)

        # Sums must be ones
        sums = xp.sum(y.array, axis=1)
        testing.assert_allclose(sums, xp.ones_like(sums))

        # Expectations must be equal to memllowmax's outputs
        testing.assert_allclose(
            xp.sum(y.array * x_data, axis=1), mellowmax(x_data, axis=1).array) 

def _batch_matmul(a, b, transa=False, transb=False, transout=False):
    a = a.reshape(a.shape[:2] + (-1,))
    b = b.reshape(b.shape[:2] + (-1,))
    trans_axis = (0, 2, 1)
    if transout:
        transa, transb = not transb, not transa
        a, b = b, a
    if transa:
        a = a.transpose(trans_axis)
    if transb:
        b = b.transpose(trans_axis)
    xp = cuda.get_array_module(a)
    if xp is numpy:
        ret = numpy.empty(a.shape[:2] + b.shape[2:], dtype=a.dtype)
        for i in six.moves.range(len(a)):
            ret[i] = numpy.dot(a[i], b[i])
        return ret
    return xp.matmul(a, b) 

def calc_loss(self, x, t):
        batch_predictions, _, grids = x
        self.xp = cuda.get_array_module(batch_predictions, t)

        loss = self.calc_actual_loss(batch_predictions, None, t)

        # reshape grids
        batch_size = t.shape[0]
        grids = grids[-1]
        grid_shape = grids.shape
        grids = F.reshape(grids, (-1, batch_size) + grid_shape[1:])

        grid_losses = []
        for grid in F.separate(grids, axis=0):
            with cuda.get_device_from_array(getattr(grid, 'data', grid[0].data)):
                grid_losses.append(self.calc_direction_loss(grid))

        return loss + (sum(grid_losses) / len(grid_losses)) 

def backward_gpu(self, inputs, gys):
        if not self.gpu_optim:
            return self.backward_cpu(inputs,  gys)
        xp = cuda.get_array_module(*inputs)
        x, gamma, beta = inputs
        gy, = gys
        g_beta = xp.sum(gy, axis=0, keepdims=True)
        g_gamma = xp.sum(gy*self.normalized, axis=0, keepdims=True)
        
        gy2 = gy*gamma
        gy_centered = gy2 - xp.mean(gy2, axis=1, keepdims=True)
        sc_prod = xp.sum(gy_centered * self.normalized, axis = 1, keepdims=True)
        
        H = x.shape[1]
#         ga = backprop_scale(self.inv_norm, gy_centered, self.normalized, sc_prod/H)
        ga = cp.ElementwiseKernel(
         'T inv_norm, T gy_centered, T normalized, T sc_prod',
         'T z',
          '''
              z = inv_norm *(gy_centered - normalized * (sc_prod/%f));
         '''%H,
         'backprop_scale')(self.inv_norm, gy_centered, self.normalized, sc_prod)
        
        return ga, g_gamma, g_beta 

def loss_unlabeled(forward, x, args):
    if args.method == 'vat':
        # Virtual adversarial training loss
        logit = forward(x, train=True, update_batch_stats=False)
        return loss.vat_loss(forward, loss.distance, x, epsilon=args.epsilon, xi=XI, p_logit=logit.data)
    elif args.method == 'vatent':
        # Virtual adversarial training loss + Conditional Entropy loss
        logit = forward(x, train=True, update_batch_stats=False)
        vat_loss = loss.vat_loss(forward, loss.distance, x, epsilon=args.epsilon, xi=XI, p_logit=logit.data)
        ent_y_x = loss.entropy_y_x(logit)
        return vat_loss + ent_y_x
    elif args.method == 'baseline':
        xp = cuda.get_array_module(x.data)
        return Variable(xp.array(0, dtype=xp.float32))
    else:
        raise NotImplementedError 

def vat_loss(forward, distance, x, train=True, epsilon=8.0, xi=1e-6, Ip=1, p_logit=None):
    if p_logit is None:
        p_logit = forward(x, train=train, update_batch_stats=False).data  # unchain
    else:
        assert not isinstance(p_logit, Variable)

    xp = cuda.get_array_module(x.data)
    d = xp.random.normal(size=x.shape)
    d = get_normalized_vector(d, xp) 
    for ip in range(Ip):
        x_d = Variable(x.data + xi * d.astype(xp.float32))
        p_d_logit = forward(x_d, train=train, update_batch_stats=False)
        kl_loss = distance(p_logit, p_d_logit)
        kl_loss.backward()
        d = x_d.grad
        d = d / xp.sqrt(xp.sum(d ** 2, axis=range(1, len(d.shape)), keepdims=True))
    x_adv = x + epsilon * d 
    p_adv_logit = forward(x_adv, train=train, update_batch_stats=False)
    return distance(p_logit, p_adv_logit) 

def forward(self, x, t):
        xp = cuda.get_array_module(x)
        y = self.predictor(x)
        log_softmax = F.log_softmax(y)
        # SelectItem is not supported by onnx-chainer.
        # TODO(hamaji): Support it?
        # log_prob = F.select_item(log_softmax, t)

        batch_size = chainer.Variable(xp.array(t.size, xp.float32),
                                      name='batch_size')
        self.extra_inputs = [batch_size]
        # TODO(hamaji): Currently, F.sum with axis=1 cannot be
        # backpropped properly.
        # log_prob = F.sum(log_softmax * t, axis=1)
        # return -F.sum(log_prob, axis=0) / self.batch_size
        log_prob = F.sum(log_softmax * t, axis=(0, 1))
        loss = -log_prob / batch_size
        reporter.report({'loss': loss}, self)
        if self.compute_accuracy:
            acc = accuracy.accuracy(y, xp.argmax(t, axis=1))
            reporter.report({'accuracy': acc}, self)
        loss.name = 'loss'
        return loss 

def calc_loss(self, x, t):
        batch_predictions, _, _ = x

        # concat all individual predictions and slice for each time step
        batch_predictions = F.concat([F.expand_dims(p, axis=0) for p in batch_predictions], axis=0)

        self.xp = cuda.get_array_module(batch_predictions[0], t)
        batch_size = t.shape[0]
        t = F.reshape(t, (batch_size, self.num_timesteps, -1))

        losses = []
        for predictions, labels in zip(F.separate(batch_predictions, axis=0), F.separate(t, axis=1)):
            batch_size, num_chars, num_classes = predictions.shape
            predictions = F.reshape(predictions, (batch_size * num_chars, num_classes))
            labels = F.reshape(labels, (-1,))
            losses.append(F.softmax_cross_entropy(predictions, labels))

        return sum(losses) 

def channelwise_inhibited(self, h):
        xp = cuda.get_array_module(h.data)
        num = h.data.shape[0]

        h = F.split_axis(h, 3, 1)
        c = F.reshape(h[self.c], (num, 16, 16))
        z = Variable(xp.zeros_like(c.data), 'AUTO')
        c = F.batch_matmul(c, z)
        c = F.reshape(c, (num, 1, 16, 16))
        hs = []
        for i, s in enumerate(h):
            if i == self.c:
                hs.append(c)
            else:
                hs.append(s)
        return F.concat(hs, 1) 

def channelwise_inhibited(self, h):
        self.c = random.randint(0, 2)
        xp = cuda.get_array_module(h.data)
        num = h.data.shape[0]

        h = F.split_axis(h, 3, 1)
        c = F.reshape(h[self.c], (num, 16, 16))
        z = Variable(xp.zeros_like(c.data), 'AUTO')
        c = F.batch_matmul(c, z)
        c = F.reshape(c, (num, 1, 16, 16))
        hs = []
        for i, s in enumerate(h):
            if i == self.c:
                hs.append(c)
            else:
                hs.append(s)
        return F.concat(hs, 1) 

def listmle(x, t):
    """
    The ListMLE loss as in Xia et al (2008), Listwise Approach to Learning to
    Rank - Theory and Algorithm.

    :param x: The activation of the previous layer
    :param t: The target labels
    :return: The loss
    """

    # Get the ground truth by sorting activations by the relevance labels
    xp = cuda.get_array_module(t)
    t_hat = t[:, 0]
    x_hat = x[xp.flip(xp.argsort(t_hat), axis=0)]

    # Compute MLE loss
    final = logcumsumexp(x_hat)
    return F.sum(final - x_hat) 

def _pl_sample(t, α):
    """
    Sample from the plackett luce distribution directly

    :param t: The target labels
    :return: A random permutation from the plackett-luce distribution
             parameterized by the target labels
    """
    xp = cuda.get_array_module(t)
    t = t[:, 0]

    probs = xp.exp(t * α)
    probs /= xp.sum(probs)

    # Use CPU-based numpy implementation, because cupy.random.choice with
    # replace=False does not work
    probs = cuda.to_cpu(probs)
    result = np.random.choice(probs.shape[0], probs.shape[0], replace=False,
                              p=probs)
    return xp.array(result, copy=False) 

def bound_by_tanh(x, low, high):
    """Bound a given value into [low, high] by tanh.

    Args:
        x (chainer.Variable): value to bound
        low (numpy.ndarray): lower bound
        high (numpy.ndarray): upper bound
    Returns: chainer.Variable
    """
    assert isinstance(x, chainer.Variable)
    assert low is not None
    assert high is not None
    xp = cuda.get_array_module(x.array)
    x_scale = (high - low) / 2
    x_scale = xp.expand_dims(xp.asarray(x_scale), axis=0)
    x_mean = (high + low) / 2
    x_mean = xp.expand_dims(xp.asarray(x_mean), axis=0)
    return F.tanh(x) * x_scale + x_mean 

def backward_log_softmax(self, x, y, gy):
        if cuda.cudnn_enabled:
            cudnn = cuda.cudnn
            libcudnn = cuda.cuda.cudnn
            _algorithm = libcudnn.CUDNN_SOFTMAX_LOG
            _mode = libcudnn.CUDNN_SOFTMAX_MODE_CHANNEL

        xp = cuda.get_array_module(x)
        if xp is not numpy and chainer.should_use_cudnn('>=auto', 3000):
            oz_dtype = 'd' if x.dtype == 'd' else 'f'
            one = numpy.array(1, dtype=oz_dtype).ctypes
            zero = numpy.array(0, dtype=oz_dtype).ctypes
            handle = cudnn.get_handle()
            gx = xp.empty(x.shape, dtype=x.dtype)
            gx_cube = gx.reshape(gx.shape[:2] + (-1, 1))
            desc = cudnn.create_tensor_descriptor(gx_cube)
            libcudnn.softmaxBackward(
                handle, _algorithm, _mode, one.data, desc.value,
                y.data.ptr, desc.value, gy.data.ptr, zero.data,
                desc.value, gx.data.ptr)
        else:
            gx = gy - xp.exp(y) * gy.sum(axis=1, keepdims=True)

        return gx 

def check_forward(self, x_data):
        xp = cuda.get_array_module(x_data)
        y = mellowmax(x_data, axis=self.axis, omega=self.omega)
        self.assertEqual(y.array.dtype, self.dtype)

        x_min = xp.min(x_data, axis=self.axis)
        x_max = xp.max(x_data, axis=self.axis)
        x_mean = xp.mean(x_data, axis=self.axis)
        print('x_min', x_min)
        print('y.array', y.array)

        # min <= mellowmax <= max
        eps = 1e-5
        self.assertTrue(xp.all(x_min <= y.array + eps))
        self.assertTrue(xp.all(x_max >= y.array - eps))

        # omega > 0 -> mellowmax is more like max
        if self.omega > 0:
            self.assertTrue(xp.all(x_mean <= y.array + eps))
        # omega < 0 -> mellowmax is more like min
        if self.omega < 0:
            self.assertTrue(xp.all(x_mean >= y.array - eps)) 

def dc_loss(embedding, label):
    """
    Deep clustering loss function.

    Args:
      embedding: (T,D)-shaped activation values
      label: (T,C)-shaped labels
    return:
      (1,)-shaped squared flobenius norm of the difference
      between embedding and label affinity matrices
    """
    xp = cuda.get_array_module(label)
    b = xp.zeros((label.shape[0], 2**label.shape[1]))
    b[np.arange(label.shape[0]),
      [int(''.join(str(x) for x in t), base=2) for t in label.data]] = 1

    label_f = chainer.Variable(b.astype(np.float32))
    loss = F.sum(F.square(F.matmul(embedding, embedding, True, False))) \
        + F.sum(F.square(F.matmul(label_f, label_f, True, False))) \
        - 2 * F.sum(F.square(F.matmul(embedding, label_f, True, False)))
    return loss 

def forward(self, inputs):
        xp = cuda.get_array_module(*inputs)
        x, t = inputs
        t = t[:, None]
        positive, unlabeled = t == self.positive, t == self.unlabeled
        n_positive, n_unlabeled = max([1., xp.sum(positive)]), max([1., xp.sum(unlabeled)])
        self.x_in = Variable(x)
        y_positive = self.loss_func(self.x_in)
        y_unlabeled = self.loss_func(-self.x_in)
        positive_risk = F.sum(self.prior * positive / n_positive * y_positive)
        negative_risk = F.sum((unlabeled / n_unlabeled - self.prior * positive / n_positive) * y_unlabeled)
        objective = positive_risk + negative_risk
        if self.nnPU:
            if negative_risk.data < -self.beta:
                objective = positive_risk - self.beta
                self.x_out = -self.gamma * negative_risk
            else:
                self.x_out = objective
        else:
            self.x_out = objective
        self.loss = xp.array(objective.data, dtype=self.x_out.data.dtype)
        return self.loss, 

def init_conv(array):
    xp = cuda.get_array_module(array)
    array[...] = xp.random.normal(loc=0.0, scale=0.02, size=array.shape) 

def at_loss(forward, x, y, train=True, epsilon=8.0):
    ce = cross_entropy(forward(x, train=train, update_batch_stats=False), y)
    ce.backward()
    d = x.grad
    xp = cuda.get_array_module(x.data)
    d = get_normalized_vector(d, xp) 
    x_adv = x + epsilon * d 
    return cross_entropy(forward(x_adv, train=train, update_batch_stats=False), y) 

def add_noise(h, sigma=0.2):
    xp = cuda.get_array_module(h.data)
    if chainer.config.train:
        return h + sigma * xp.random.randn(*h.shape)
    else:
        return h 

def forward(self, x):
        self.retain_outputs((0,))
        xp = cuda.get_array_module(*x)
        return utils.force_array(1 / xp.sqrt(x[0], dtype=x[0].dtype)), 

def init_bn(array):
    xp = cuda.get_array_module(array)
    array[...] = xp.random.normal(loc=1.0, scale=0.02, size=array.shape) 

def kl_binary(p_logit, q_logit):
    if isinstance(p_logit, chainer.Variable):
        xp = cuda.get_array_module(p_logit.data)
    else:
        xp = cuda.get_array_module(p_logit)
    p_logit = F.concat([p_logit, xp.zeros(p_logit.shape, xp.float32)], 1)
    q_logit = F.concat([q_logit, xp.zeros(q_logit.shape, xp.float32)], 1)
    return kl_categorical(p_logit, q_logit) 

def add_noise(h, sigma=0.2):
    xp = cuda.get_array_module(h.data)
    if not chainer.config.train:
        return h
    else:
        return h + sigma * xp.random.randn(*h.data.shape)


# differentiable backward functions 

def kl_categorical(p_logit, q_logit):
    if isinstance(p_logit, chainer.Variable):
        xp = cuda.get_array_module(p_logit.data)
    else:
        xp = cuda.get_array_module(p_logit)
    p = F.softmax(p_logit)
    _kl = F.sum(p * (F.log_softmax(p_logit) - F.log_softmax(q_logit)), 1)
    return F.sum(_kl) / xp.prod(xp.array(_kl.shape)) 

def source_pad_concat_convert(x_seqs, device, eos_id=0):
    x_block = convert.concat_examples(x_seqs, device, padding=-1)
    xp = cuda.get_array_module(x_block)

    # add eos
    x_block = xp.pad(x_block, ((0, 0), (0, 1)),
                     'constant', constant_values=-1)
    for i_batch, seq in enumerate(x_seqs):
        x_block[i_batch, len(seq)] = eos_id
    return x_block 

def forward(self, x):
        self.retain_outputs((0,))
        xp = cuda.get_array_module(*x)
        return utils.force_array(xp.sqrt(x[0], dtype=x[0].dtype)), 

def remove_sentence_boundaries(tensor,
                               mask):
    """
    Remove begin/end of sentence embeddings from the batch of sentences.
    Given a batch of sentences with size ``(batch_size, timesteps, dim)``
    this returns a tensor of shape ``(batch_size, timesteps - 2, dim)`` after removing
    the beginning and end sentence markers.  The sentences are assumed to be padded on the right,
    with the beginning of each sentence assumed to occur at index 0 (i.e., ``mask[:, 0]`` is assumed
    to be 1).
    Returns both the new tensor and updated mask.
    This function is the inverse of ``add_sentence_boundary_token_ids``.
    Parameters
    ----------
    tensor : ``torch.Tensor``
        A tensor of shape ``(batch_size, timesteps, dim)``
    mask : ``torch.Tensor``
         A tensor of shape ``(batch_size, timesteps)``
    Returns
    -------
    tensor_without_boundary_tokens : ``torch.Tensor``
        The tensor after removing the boundary tokens of shape ``(batch_size, timesteps - 2, dim)``
    new_mask : ``torch.Tensor``
        The new mask for the tensor of shape ``(batch_size, timesteps - 2)``.
    """
    xp = cuda.get_array_module(mask)

    # TODO: matthewp, profile this transfer
    sequence_lengths = mask.sum(axis=1)
    tensor_shape = list(tensor.array.shape)
    new_shape = list(tensor_shape)
    new_shape[1] = tensor_shape[1] - 2
    tensor_without_boundary_tokens = xp.zeros(new_shape, 'f')
    new_mask = xp.zeros((new_shape[0], new_shape[1]), 'i')
    for i, j in enumerate(sequence_lengths):
        if j > 2:
            tensor_without_boundary_tokens[i, :(j - 2), :] = \
                tensor.array[i, 1:(j - 1), :]
            new_mask[i, :(j - 2)] = 1

    return tensor_without_boundary_tokens, new_mask 

def init_state(self, param, state):
        xp = cuda.get_array_module(param.data)
        with cuda.get_device(param.data):
            state['v'] = xp.zeros_like(param.data) 

def forward(self, x):
        xp = cuda.get_array_module(*x)
        n, c, h_i, w_i = x[0].shape
        size_i = np.asarray(x[0].shape[2:4])
        size_o = self.size
        h_o, w_o = size_o
        y = xp.zeros(shape=(n,c,h_o,w_o), dtype=np.float32)

        # [-1, 1]^2 -> [0, h_i]x[0, w_i]
        center = 0.5 * (self.center+1) * (size_i+1)

        # topleft: np.array[batch, [top, left]]
        topleft = center - 0.5*size_o
        topleft = np.round(topleft).astype(np.int32)

        tl_o = np.maximum(topleft, 0)
        br_o = np.minimum(topleft+size_o, size_i)

        tl_i = tl_o - topleft
        br_i = br_o - topleft

        for k in range(n):
            if (br_i[k,0] - tl_i[k,0]) < 0 or (br_i[k,1] - tl_i[k,1]) < 0:
                continue
            y[k,:,tl_i[k,0]:br_i[k,0],tl_i[k,1]:br_i[k,1]] \
                += x[0][k,:,tl_o[k,0]:br_o[k,0],tl_o[k,1]:br_o[k,1]]

        return y,

    # do not backward (always return 0)