Mpi 在pytorch数据并行模式下，如何使用全局张量？

在这个例子中，我希望z_proto对于不同的gpu是全局的。然而，在数据并行模式下，它也被分成不同的GPU。如何解决这样的问题？多谢各位

class SequencePrototypeTokenClassification(nn.Module):
    def __init__(self,seq_model, label_num):
        super(SequencePrototypeTokenClassification, self).__init__()
        self.seq_model = seq_model
        self.label_num = label_num

    def forward(self, input_ids, token_type_ids, attention_mask, labels, z_proto, n_query, target_inds):
        z, _ = self.seq_model(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False)
        z_dim = z.size(-1)
        zq = z.squeeze().view(-1, z_dim)
        dists = euclidean_dist(zq, z_proto)
        log_p_y = F.log_softmax(-dists, dim=1).view(-1, self.label_num)
        loss_val = -log_p_y.gather(1, self.target_inds).squeeze().view(-1).mean()
        _, y_hat = log_p_y.max(1)

        return loss_val, y_hat

---

根据上面的代码，

z_proto

似乎是forward函数的参数之一，而不是模型的一部分。因此，只需将其存储在主GPU上的

张量
中，就可以使其在GPU上具有相同的值

编辑
基于，似乎
DataParallel
将前向传递函数的所有输入分割到GPU上。通过将其作为类变量存储在模型对象本身中，您可以绕过它。如果forward函数不是静态变量，则可以在调用该函数之前更新该值

class SequencePrototypeTokenClassification（nn.Module）：
定义初始值（自、序列型号、标签编号）：
...
self.z_proto=无
...
...
#训练回路
...
model.z_proto=值
model.forward（）
...
事实证明，
DataParallel
只会复制
nn.Module
的
nn.Parameter
。因此，我在模块中随机初始化了一个名为
z_proto
的
nn.参数
，并将张量
z_proto
的值复制到参数中。然后将参数复制到4个GPU中。
我将
z_proto
放在主GPU上。但是
DataParallel
中的
replicas=self.replicate（self.module，self.device_id[：len（inputs）]）
会将
z_proto
拆分到4个GPU上。嗨，我尝试过这种方法
proto_model.module.z_proto=z_proto
但它得到了运行时错误：参数位于不同的GPU上，位于/pytorch/aten/src/THC/generated/。/generic/THCTensorMathPointwise.cu:272
，这很奇怪。根据，pytorch仅在转发呼叫期间进行拆分，并在下一行之前将其合并回来。您确定已将
z_proto
更新置于
forward（）
函数之外吗？你在
\uuu init\uuu
函数中初始化了
z_proto
了吗？是的，我做的事情和你的几乎完全一样。我用
None
初始化了
z_proto
，并在
forward
之前更新它，但只使用GPU 0。我原以为
z_proto
会与
模块一起复制，但事实并非如此。你能将完整的代码与培训循环一起分享吗？