一.参考链接
二.代码
import torch.nn as nn
class AdaptiveInput(nn.Module):
"""
This implementation and the above description are heavily cited from the softmax counterpart from
https://pytorch.org/docs/stable/_modules/torch/nn/modules/adaptive.html
"""
def __init__(self, in_features, n_classes, cutoffs=None,div_value=4., head_bias=False):
"""
:param in_features:embeding的维度论文中的参数d
:param n_classes:词汇表长度
:param cutoffs: 是一个list,控制词汇表的划分,例如:cutoffs = [10, 100, 1000]
表示将词汇表划分成五个clusters:[0-10]、[11-100]、[101-1000]和[1001-最后]
:param div_value:就是论文中的参数 k
:param head_bias:
"""
super(AdaptiveInput, self).__init__()
if not cutoffs:
cutoffs = [10000, 60000, 190000]
cutoffs = list(cutoffs)
if (cutoffs != sorted(cutoffs)) \
or (min(cutoffs) <= 0) \
or (max(cutoffs) >= (n_classes - 1)) \
or (len(set(cutoffs)) != len(cutoffs)) \
or any([int(c) != c for c in cutoffs]):
raise ValueError("cutoffs should be a sequence of unique, positive "
"integers sorted in an increasing order, where "
"each value is between 1 and n_classes-1")
self.in_features = in_features
self.n_classes = n_classes
self.cutoffs = cutoffs + [n_classes]
self.div_value = div_value
self.head_bias = head_bias
self.n_clusters = len(self.cutoffs) - 1
self.head_size = self.cutoffs[0]
self.head = nn.Sequential(nn.Embedding(self.head_size, self.in_features),
nn.Linear(self.in_features, self.in_features, bias=self.head_bias))
self.tail = nn.ModuleList()
for i in range(self.n_clusters):
hsz = int(self.in_features // (self.div_value ** (i + 1)))
osz = self.cutoffs[i + 1] - self.cutoffs[i]
projection = nn.Sequential(
nn.Embedding(osz, hsz),
nn.Linear(hsz, self.in_features, bias=False),
)
self.tail.append(projection)
def forward(self, input):
"""
:param input: 一个句子list,中间元素是句子中词在词汇表中的编号
"""
used_rows = 0
input_size = list(input.size())
output = input.new_zeros(input_size + [self.in_features]).float()
cutoff_values = [0] + self.cutoffs
for i in range(len(cutoff_values) - 1):
low_idx = cutoff_values[i]
high_idx = cutoff_values[i + 1]
input_mask = (input >= low_idx) & (input < high_idx)
row_indices = input_mask.nonzero().squeeze()
if row_indices.numel() == 0:
continue
out = self.head(input[input_mask] - low_idx) if i == 0 else self.tail[i - 1](input[input_mask] - low_idx)
output.index_copy_(0, row_indices, out)
used_rows += row_indices.numel()
if used_rows != input_size[0]:
raise RuntimeError("Target values should be in [0, {}], "
"but values in range [{}, {}] "
"were found. ".format(self.n_classes - 1,
input.min().item(),
input.max().item()))
return output
import torch
x = torch.arange(0,100).long()
inp = AdaptiveInput(128, 100, cutoffs=[4,8,16])
print(inp(x))