This repository has been archived by the owner on Feb 7, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 2k
/
uniform_sampling.py
99 lines (86 loc) · 3.03 KB
/
uniform_sampling.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
## @package uniform_sampling
# Module caffe2.python.layers.uniform_sampling
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import numpy as np
from caffe2.python import core, schema
from caffe2.python.layers.layers import LayerParameter, ModelLayer
class UniformSampling(ModelLayer):
"""
Uniform sampling `num_samples - len(input_record)` unique elements from the
range [0, num_elements). `samples` is the concatenation of input_record and
the samples. input_record is expected to be unique.
"""
def __init__(
self,
model,
input_record,
num_samples,
num_elements,
name='uniform_sampling',
**kwargs
):
super(UniformSampling, self).__init__(
model, name, input_record, **kwargs
)
assert num_elements > 0
assert isinstance(input_record, schema.Scalar)
self.num_elements = num_elements
self.num_samples = model.net.NextScopedBlob(name + "_num_samples")
self.params.append(
LayerParameter(
parameter=self.num_samples,
initializer=core.CreateOperator(
"GivenTensorInt64Fill",
[],
self.num_samples,
shape=(1, ),
values=[num_samples],
),
optimizer=model.NoOptim,
)
)
self.sampling_prob = model.net.NextScopedBlob(name + "_prob")
self.params.append(
LayerParameter(
parameter=self.sampling_prob,
initializer=core.CreateOperator(
"ConstantFill",
[],
self.sampling_prob,
shape=(num_samples, ),
value=float(num_samples) / num_elements,
dtype=core.DataType.FLOAT
),
optimizer=model.NoOptim,
)
)
self.output_schema = schema.Struct(
(
'samples', schema.Scalar(
np.int32, model.net.NextScopedBlob(name + "_samples")
)
),
('sampling_prob', schema.Scalar(np.float32, self.sampling_prob)),
)
def add_ops(self, net):
net.StopGradient(self.sampling_prob, self.sampling_prob)
shape = net.Shape([self.input_record()], net.NextScopedBlob("shape"))
shape = net.Sub([self.num_samples, shape], shape)
samples = net.UniqueUniformFill(
[shape, self.input_record()],
net.NextScopedBlob("samples"),
min=0,
max=self.num_elements - 1,
input_as_shape=True
)
net.Concat(
[self.input_record(), samples],
[self.output_schema.samples(), net.NextScopedBlob("split_info")],
axis=0
)
net.StopGradient(
self.output_schema.samples(), self.output_schema.samples()
)