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

Question

在这个例子中，我希望 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

Answer 1

根据您上面的代码，z_proto 似乎是前向函数的参数之一，而不是模型的一部分。因此，只需将它存储在主 GPU 上的 tensor 中，就可以使其在 GPU 之间具有相同的值。

编辑

基于the documentation，似乎DataParallel 将所有输入拆分为跨 GPU 的前向传递函数。一种可以规避它的方法是将其存储为模型对象本身内的 class 变量。如果它不是静态变量，则可以在调用前向函数之前更新该值。

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


#Training loop
    ...
    model.z_proto = value
    model.forward()
    ...

Answer 2

原来 DataParallel 只会复制 nn.Module 的 nn.Parameter。所以我在模块中随机初始化了一个名为z_proto的nn.Parameter，并将张量z_proto的值复制到参数中。然后将参数复制到4个GPU中。

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

In pytorch data parallel mode, how to use the global tensor?

mpi

pytorch

编辑