import torch
import torch.nn as nn
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1,2,3"
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.counter = torch.tensor(0)
def forward(self, x):
self.counter += 1
return x
# 假设有输入数据x
model = MyModel()
# 使用四个GPU进行数据并行训练
model = nn.DataParallel(model, device_ids=[0, 1,2,3])
x = torch.rand((128,))
x = x.to("cuda:0")
# 第一次调用forward
output1 = model(x)
# 获取counter的值
counter_value = model.module.counter.item()
print("Counter value:", counter_value)
输出为4