mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import numpy as np
|
|
import unittest
|
|
from caffe2.proto import caffe2_pb2
|
|
from caffe2.python import core, workspace, data_parallel_model
|
|
from caffe2.python.test_util import TestCase
|
|
|
|
|
|
@unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
|
|
@unittest.skipIf(workspace.NumCudaDevices() < 2, "Need at least 2 GPUs.")
|
|
class GPUDataParallelModelTest(TestCase):
|
|
def test(self):
|
|
gpu_devices = [0, 1] # gpu ids
|
|
perfect_model = np.array([2, 6, 5, 0, 1]).astype(np.float32)
|
|
np.random.seed(123)
|
|
data = np.random.randint(
|
|
2, size=(50, perfect_model.size)
|
|
).astype(np.float32)
|
|
label = np.dot(data, perfect_model)[:, np.newaxis]
|
|
|
|
model = data_parallel_model.GPUDataParallelModel(
|
|
gpu_devices, order="NHWC", name="fake")
|
|
|
|
fc = model.FC("data", "fc", perfect_model.size, 1,
|
|
("ConstantFill", {}), ("ConstantFill", {}), axis=0)
|
|
sq = model.SquaredL2Distance([fc, "label"], "sq")
|
|
loss = model.AveragedLoss(sq, "loss")
|
|
model.AddGradientOperators([loss])
|
|
model.SGD(-0.1)
|
|
model.RunAllOnGPU()
|
|
|
|
for gpu_id in gpu_devices:
|
|
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, gpu_id)):
|
|
workspace.FeedBlob(
|
|
"gpu_{}/data".format(gpu_id), data[0])
|
|
workspace.FeedBlob(
|
|
"gpu_{}/label".format(gpu_id), label[0])
|
|
|
|
workspace.RunNetOnce(model.param_init_net)
|
|
workspace.CreateNet(model.net)
|
|
|
|
for i in range(2000):
|
|
idx = np.random.randint(data.shape[0])
|
|
for gpu_id in gpu_devices:
|
|
device = core.DeviceOption(caffe2_pb2.CUDA, gpu_id)
|
|
with core.DeviceScope(device):
|
|
workspace.FeedBlob(
|
|
"gpu_{}/data".format(gpu_id), data[idx])
|
|
workspace.FeedBlob(
|
|
"gpu_{}/label".format(gpu_id), label[idx])
|
|
workspace.RunNet(model.net)
|
|
|
|
for gpu_id in gpu_devices:
|
|
np.testing.assert_allclose(
|
|
perfect_model[np.newaxis, :],
|
|
workspace.FetchBlob("gpu_{}/fc_w".format(gpu_id)),
|
|
atol=1e-2)
|