Method update

python/dgl/optim/pytorch/sparse_optim.py:814–938

Update embeddings in a sparse manner Sparse embeddings are updated in mini batches. We maintain gradient states for each embedding so they can be updated separately. Parameters ---------- idx : tensor Index of the embeddings to be updated.

(self, idx, grad, emb)

            emb.set_optm_state(state)

    def update(self, idx, grad, emb):
        """Update embeddings in a sparse manner
        Sparse embeddings are updated in mini batches. We maintain gradient states for
        each embedding so they can be updated separately.

        Parameters
        ----------
        idx : tensor
            Index of the embeddings to be updated.
        grad : tensor
            Gradient of each embedding.
        emb : dgl.nn.NodeEmbedding
            Sparse embedding to update.
        """
        with th.no_grad():
            state_step, state_mem, state_power = emb.optm_state
            exec_dtype = grad.dtype
            exec_dev = grad.device
            state_dev = state_step.device

            # whether or not we need to transfer data from the GPU to the CPU
            # while updating the weights
            is_d2h = state_dev.type == "cpu" and exec_dev.type == "cuda"

            # only perform async copies cpu -> gpu, or gpu-> gpu, but block
            # when copying to the cpu, so as to ensure the copy is finished
            # before operating on the data on the cpu
            state_block = is_d2h

            if self._is_using_uva[emb.name] is None and is_d2h:
                # we should use UVA going forward
                self._setup_uva(emb.name, state_mem, state_power)
            elif self._is_using_uva[emb.name] is None:
                # we shouldn't use UVA going forward
                self._is_using_uva[emb.name] = False

            use_uva = self._is_using_uva[emb.name]

            beta1 = self._beta1
            beta2 = self._beta2
            eps = self._eps

            clr = self._lr
            # There can be duplicated indices due to sampling.
            # Thus unique them here and average the gradient here.
            grad_indices, inverse, cnt = th.unique(
                idx, return_inverse=True, return_counts=True
            )
            state_idx = grad_indices.to(state_dev)
            state_step[state_idx] += 1
            state_step = state_step[state_idx].to(exec_dev)

            if use_uva:
                orig_mem = gather_pinned_tensor_rows(state_mem, grad_indices)
                orig_power = gather_pinned_tensor_rows(
                    state_power, grad_indices
                )
            else: