Method setup

python/dgl/optim/pytorch/sparse_optim.py:555–591

(self, params)

        self.setup(self._params)

    def setup(self, params):
        # We need to register a state sum for each embedding in the kvstore.
        for emb in params:
            assert isinstance(
                emb, NodeEmbedding
            ), "SparseAdagrad only supports dgl.nn.NodeEmbedding"

            emb_name = emb.name
            if th.device(emb.weight.device) == th.device("cpu"):
                # if our embedding is on the CPU, our state also has to be
                if self._rank < 0:
                    state = th.empty(
                        emb.weight.shape,
                        dtype=th.float32,
                        device=th.device("cpu"),
                    ).zero_()
                elif self._rank == 0:
                    state = create_shared_mem_array(
                        emb_name + "_state", emb.weight.shape, th.float32
                    ).zero_()

                    if self._world_size > 1:
                        emb.store.set(emb_name + "_opt", emb_name)
                elif self._rank > 0:
                    # receive
                    emb.store.wait([emb_name + "_opt"])
                    state = get_shared_mem_array(
                        emb_name + "_state", emb.weight.shape, th.float32
                    )
            else:
                # distributed state on on gpu
                state = th.empty(
                    emb.weight.shape,
                    dtype=th.float32,
                    device=emb.weight.device,
                ).zero_()
            emb.set_optm_state((state,))

    def update(self, idx, grad, emb):
        """Update embeddings in a sparse manner