-
Notifications
You must be signed in to change notification settings - Fork 1.4k
/
black_out.py
61 lines (44 loc) · 1.78 KB
/
black_out.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
import numpy
from chainer.functions.loss import black_out
from chainer import link
from chainer.utils import walker_alias
from chainer import variable
class BlackOut(link.Link):
"""BlackOut loss layer.
.. seealso:: :func:`~chainer.functions.black_out` for more detail.
Args:
in_size (int): Dimension of input vectors.
counts (int list): Number of each identifiers.
sample_size (int): Number of negative samples.
Attributes:
W (~chainer.Parameter): Weight parameter matrix.
"""
sample_data = None
def __init__(self, in_size, counts, sample_size):
super(BlackOut, self).__init__()
vocab_size = len(counts)
p = numpy.array(counts, dtype=numpy.float32)
self.sampler = walker_alias.WalkerAlias(p)
self.sample_size = sample_size
with self.init_scope():
self.W = variable.Parameter(shape=(vocab_size, in_size))
def device_resident_accept(self, visitor):
super(BlackOut, self).device_resident_accept(visitor)
self.sampler.device_resident_accept(visitor)
def forward(self, x, t):
"""Computes the loss value for given input and ground truth labels.
Args:
x (~chainer.Variable): Input of the weight matrix multiplication.
t (~chainer.Variable): Batch of ground truth labels.
Returns:
~chainer.Variable: Loss value.
"""
batch_size = x.shape[0]
if self.sample_data is not None:
# for test
sample_data = self.sample_data
else:
shape = (batch_size, self.sample_size)
sample_data = self.sampler.sample(shape)
samples = variable.Variable(sample_data, requires_grad=False)
return black_out.black_out(x, t, self.W, samples)