How can I modify the value of DistTensor?

[xcwanAndy]

Hi, all.

I found that it’s not possible for me to change the value of a DistTensor object. See below:

logger.debug("Before: {}".format(
        torch.sum(g.ndata['features'][local_nodes][0])))
g.ndata['features'][local_nodes][0] = torch.zeros(
     g.ndata['features'][local_nodes][0].shape)
 logger.debug("After: {}".format(
     torch.sum(g.ndata['features'][local_nodes][0])))

After running this program, I found the value unchanged. But I thought the source code of DistTensor supports assigning tensor values via pushing the new value to kvstore:

    def __setitem__(self, idx, val):
        idx = utils.toindex(idx)
        idx = idx.tousertensor()
        # TODO(zhengda) how do we want to support broadcast (e.g., G.ndata['h'][idx] = 1).
        self.kvstore.push(name=self._name, id_tensor=idx, data_tensor=val)

I wonder if it’s because DGL by default store caches after the first read, and afterwards only reads values from the local cache? Or is the DistTensor by default is read-only?

Besides, as I only need to change the feature values of local nodes, do I have any approach to achieve it?

[xcwanAndy]

@VoVAllen Would you please give me some suggestions?

[VoVAllen]

There shouldn’t be any consistency(cache) problem, since the tensor only has one storage in shared memory on each machine. What’s [0] in g.ndata['features'][local_nodes][0]? It seems you are modifying a local tensor instead of DistTensor

[xcwanAndy]

I wrote [0] because DistTensor has ntypes. As reddit only has one type, so I set 0 by default.

I want to modify the local tensor at each worker, so that when other workers run dgl.distributed.sample_neighbors as their collate_fn, they can obtain the newly updated tensors.

Would you please give me some advice to help satisfy the above requirement?

[VoVAllen]

The __setitem__ function should work in your case. Did you meet any problem? Could you try use this function explicitly?