mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
6a4182eead
Summary: CUDA version of weighted sampling operator; minor changes for CPU version Reviewed By: asaadaldien Differential Revision: D6106668 fbshipit-source-id: 42d7607bd845a4a39cf5b89d7476904cb5928431
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
import numpy as np
|
|
|
|
from hypothesis import given
|
|
import hypothesis.strategies as st
|
|
|
|
from caffe2.python import core
|
|
from caffe2.python import workspace
|
|
import caffe2.python.hypothesis_test_util as hu
|
|
|
|
|
|
class TestWeightedSample(hu.HypothesisTestCase):
|
|
@given(
|
|
batch=st.integers(min_value=0, max_value=128),
|
|
weights_len=st.integers(min_value=0, max_value=128),
|
|
**hu.gcs
|
|
)
|
|
def test_weighted_sample(self, batch, weights_len, gc, dc):
|
|
op = core.CreateOperator(
|
|
"WeightedSample",
|
|
["weights"],
|
|
["indices"],
|
|
)
|
|
|
|
weights = np.zeros((batch, weights_len))
|
|
rand_indices = []
|
|
if batch > 0 and weights_len > 0:
|
|
for i in range(batch):
|
|
rand_tmp = np.random.randint(0, weights_len)
|
|
rand_indices.append(rand_tmp)
|
|
weights[i, rand_tmp] = 1.0
|
|
|
|
rand_indices = np.array(rand_indices, dtype=np.float32)
|
|
workspace.FeedBlob("weights", weights.astype(np.float32))
|
|
workspace.RunOperatorOnce(op)
|
|
result = workspace.FetchBlob("indices")
|
|
|
|
if batch > 0 and weights_len > 0:
|
|
for i in range(batch):
|
|
np.testing.assert_allclose(rand_indices[i], result[i])
|
|
else:
|
|
np.testing.assert_allclose(rand_indices, result)
|
|
|
|
self.assertDeviceChecks(dc, op, [weights.astype(np.float32)], [0])
|
|
|
|
if __name__ == "__main__":
|
|
import unittest
|
|
unittest.main()
|