Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
曾经的那一瞬间
Models
提交
1d8c6485
M
Models
项目概览
曾经的那一瞬间
/
Models
大约 1 年 前同步成功
通知
1
Star
0
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
M
Models
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
1d8c6485
编写于
10月 19, 2022
作者:
H
Hongkun Yu
提交者:
A. Unique TensorFlower
10月 19, 2022
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
Internal change
PiperOrigin-RevId: 482230148
上级
c31b02b2
变更
1
隐藏空白更改
内联
并排
Showing
1 changed file
with
67 addition
and
68 deletion
+67
-68
official/nlp/modeling/layers/moe_test.py
official/nlp/modeling/layers/moe_test.py
+67
-68
未找到文件。
official/nlp/modeling/layers/moe_test.py
浏览文件 @
1d8c6485
...
...
@@ -14,30 +14,28 @@
"""Tests for moe.py."""
import
ml_collections
import
numpy
as
np
import
tensorflow
as
tf
from
official.nlp.modeling.layers
import
moe
def
small_config
()
->
ml_collections
.
ConfigDict
:
def
small_config
():
"""Creates a small model config that can be used by all tests."""
config
=
ml_collections
.
ConfigDict
()
config
.
d_ff
=
32
config
.
dropout_rate
=
0.1
config
.
num_experts
=
2
config
.
expert_d_ff
=
33
config
.
expert_dropout_rate
=
0.1
config
.
jitter_noise
=
0.1
config
.
train_capacity_factor
=
1.0
config
.
eval_capacity_factor
=
1.0
config
.
min_expert_capacity
=
1
config
.
max_group_size
=
9
config
.
backbone_d_ff
=
13
config
=
{}
config
[
'd_ff'
]
=
32
config
[
'dropout_rate'
]
=
0.1
config
[
'num_experts'
]
=
2
config
[
'expert_d_ff'
]
=
33
config
[
'expert_dropout_rate'
]
=
0.1
config
[
'jitter_noise'
]
=
0.1
config
[
'train_capacity_factor'
]
=
1.0
config
[
'eval_capacity_factor'
]
=
1.0
config
[
'min_expert_capacity'
]
=
1
config
[
'max_group_size'
]
=
9
config
[
'backbone_d_ff'
]
=
13
return
config
...
...
@@ -124,58 +122,61 @@ class MoeTest(tf.test.TestCase):
out_ca
=
router_mask
.
combine_array
.
numpy
()
out_ca
=
np
.
dot
(
out_ca
,
np
.
ones
((
expert_capacity
,)))
expected_combine_array
=
np
.
array
(
[[[
0.66
,
0.0
,
0.0
],
[
0.42
,
0.42
,
0.16
],
[
0.0
,
0.33
,
0.33
]],
[[
0.33
,
0.33
,
0.0
],
[
0.16
,
0.42
,
0.42
],
[
0.0
,
0.0
,
0.66
]]])
expected_combine_array
=
np
.
array
([[[
0.66
,
0.0
,
0.0
],
[
0.42
,
0.42
,
0.16
],
[
0.0
,
0.33
,
0.33
]],
[[
0.33
,
0.33
,
0.0
],
[
0.16
,
0.42
,
0.42
],
[
0.0
,
0.0
,
0.66
]]])
self
.
assertAllClose
(
expected_combine_array
,
out_ca
,
atol
=
1e-2
)
def
test_feed_forward_shape_and_vars
(
self
):
config
=
small_config
()
layer
=
moe
.
FeedForward
(
d_ff
=
config
.
d_ff
,
dropout_rate
=
config
.
dropout_rate
)
layer
=
moe
.
FeedForward
(
d_ff
=
config
[
'd_ff'
],
dropout_rate
=
config
[
'dropout_rate'
])
inputs
=
make_input_ones
()
outputs
=
layer
(
inputs
)
self
.
assertAllEqual
(
tf
.
shape
(
inputs
),
tf
.
shape
(
outputs
))
var_names
=
sorted
([
v
.
name
for
v
in
layer
.
trainable_variables
])
self
.
assertAllEqual
([
'feed_forward/intermediate/bias:0'
,
'feed_forward/intermediate/kernel:0'
,
'feed_forward/output/bias:0'
,
'feed_forward/output/kernel:0'
],
var_names
)
self
.
assertAllEqual
([
'feed_forward/intermediate/bias:0'
,
'feed_forward/intermediate/kernel:0'
,
'feed_forward/output/bias:0'
,
'feed_forward/output/kernel:0'
],
var_names
)
def
test_feed_forward_manual
(
self
):
config
=
small_config
()
layer
=
moe
.
FeedForward
(
d_ff
=
config
.
d_ff
,
dropout_rate
=
config
.
dropout_rate
,
d_ff
=
config
[
'd_ff'
]
,
dropout_rate
=
config
[
'dropout_rate'
]
,
activation
=
tf
.
keras
.
activations
.
relu
,
kernel_initializer
=
tf
.
keras
.
initializers
.
get
(
'ones'
),
bias_initializer
=
tf
.
keras
.
initializers
.
get
(
'ones'
))
inputs
=
make_input_ones
(
1
,
2
,
3
)
outputs
=
layer
(
inputs
,
training
=
False
)
manual_outputs
=
tf
.
constant
([[[
129.0
,
129.0
,
129.0
],
[
129.0
,
129.0
,
129.0
]]])
manual_outputs
=
tf
.
constant
([[[
129.0
,
129.0
,
129.0
],
[
129.0
,
129.0
,
129.0
]]])
self
.
assertAllClose
(
manual_outputs
,
outputs
,
atol
=
1e-7
)
def
test_feed_forward_experts_shape_and_vars
(
self
):
config
=
small_config
()
layer
=
moe
.
FeedForwardExperts
(
num_experts
=
config
.
num_experts
,
d_ff
=
config
.
expert_d_ff
,
dropout_rate
=
config
.
expert_dropout_rate
)
num_experts
=
config
[
'num_experts'
]
,
d_ff
=
config
[
'expert_d_ff'
]
,
dropout_rate
=
config
[
'expert_dropout_rate'
]
)
inputs
=
make_experts_input_ones
()
outputs
=
layer
(
inputs
)
self
.
assertAllEqual
(
tf
.
shape
(
inputs
),
tf
.
shape
(
outputs
))
var_names
=
sorted
([
v
.
name
for
v
in
layer
.
trainable_variables
])
self
.
assertAllEqual
([
'experts/intermediate/bias:0'
,
'experts/intermediate/kernel:0'
,
'experts/output/bias:0'
,
'experts/output/kernel:0'
],
var_names
)
self
.
assertAllEqual
([
'experts/intermediate/bias:0'
,
'experts/intermediate/kernel:0'
,
'experts/output/bias:0'
,
'experts/output/kernel:0'
],
var_names
)
def
test_feed_forward_experts_manual
(
self
):
config
=
small_config
()
layer
=
moe
.
FeedForwardExperts
(
num_experts
=
1
,
d_ff
=
config
.
expert_d_ff
,
dropout_rate
=
config
.
expert_dropout_rate
,
d_ff
=
config
[
'expert_d_ff'
]
,
dropout_rate
=
config
[
'expert_dropout_rate'
]
,
activation
=
tf
.
keras
.
activations
.
relu
,
kernel_initializer
=
tf
.
keras
.
initializers
.
get
(
'ones'
),
bias_initializer
=
tf
.
keras
.
initializers
.
get
(
'ones'
))
...
...
@@ -188,38 +189,37 @@ class MoeTest(tf.test.TestCase):
def
test_moe_layer
(
self
):
config
=
small_config
()
experts
=
moe
.
FeedForwardExperts
(
num_experts
=
config
.
num_experts
,
d_ff
=
config
.
expert_d_ff
,
dropout_rate
=
config
.
expert_dropout_rate
)
num_experts
=
config
[
'num_experts'
]
,
d_ff
=
config
[
'expert_d_ff'
]
,
dropout_rate
=
config
[
'expert_dropout_rate'
]
)
router
=
moe
.
ExpertsChooseMaskedRouter
(
config
.
num_experts
,
jitter_noise
=
config
.
jitter_noise
)
config
[
'num_experts'
],
jitter_noise
=
config
[
'jitter_noise'
])
moe_layer
=
moe
.
MoeLayer
(
experts
,
router
,
train_capacity_factor
=
config
.
train_capacity_factor
,
eval_capacity_factor
=
config
.
eval_capacity_factor
,
max_group_size
=
config
.
max_group_size
,
min_expert_capacity
=
config
.
min_expert_capacity
)
train_capacity_factor
=
config
[
'train_capacity_factor'
]
,
eval_capacity_factor
=
config
[
'eval_capacity_factor'
]
,
max_group_size
=
config
[
'max_group_size'
]
,
min_expert_capacity
=
config
[
'min_expert_capacity'
]
)
inputs
=
make_input_ones
()
with
self
.
assertLogs
(
'absl'
,
level
=
'INFO'
)
as
cm
:
outputs
=
moe_layer
(
inputs
,
training
=
True
)
self
.
assertAllEqual
(
tf
.
shape
(
inputs
),
tf
.
shape
(
outputs
))
self
.
assertEqual
(
cm
.
output
,
[
(
'INFO:absl:Selected group_size=5 and num_groups=4 for input '
'num_tokens=20, max_group_size=9, num_experts=2.'
),
(
'INFO:absl:Selected expert_capacity=2 for num_experts=2 and '
'training=True.'
)])
self
.
assertEqual
(
cm
.
output
,
[(
'INFO:absl:Selected group_size=5 and num_groups=4 for input '
'num_tokens=20, max_group_size=9, num_experts=2.'
),
(
'INFO:absl:Selected expert_capacity=2 for num_experts=2 and '
'training=True.'
)])
var_names
=
sorted
([
v
.
name
for
v
in
moe_layer
.
trainable_variables
])
self
.
assertAllEqual
([
'moe/experts/intermediate/bias:0'
,
'moe/experts/intermediate/kernel:0'
,
'moe/experts/output/bias:0'
,
'moe/experts/output/kernel:0'
,
'moe/router/router_weights/bias:0'
,
'moe/router/router_weights/kernel:0'
],
var_names
)
self
.
assertAllEqual
([
'moe/experts/intermediate/bias:0'
,
'moe/experts/intermediate/kernel:0'
,
'moe/experts/output/bias:0'
,
'moe/experts/output/kernel:0'
,
'moe/router/router_weights/bias:0'
,
'moe/router/router_weights/kernel:0'
],
var_names
)
self
.
assertLen
(
moe_layer
.
losses
,
1
)
metrics
=
[
metric
.
name
for
metric
in
moe_layer
.
metrics
]
self
.
assertSetEqual
(
...
...
@@ -231,20 +231,19 @@ class MoeTest(tf.test.TestCase):
def
test_moe_layer_with_backbone
(
self
):
config
=
small_config
()
experts
=
moe
.
FeedForwardExperts
(
num_experts
=
config
.
num_experts
,
d_ff
=
config
.
expert_d_ff
,
dropout_rate
=
config
.
expert_dropout_rate
)
num_experts
=
config
[
'num_experts'
]
,
d_ff
=
config
[
'expert_d_ff'
]
,
dropout_rate
=
config
[
'expert_dropout_rate'
]
)
router
=
moe
.
ExpertsChooseMaskedRouter
(
config
.
num_experts
,
jitter_noise
=
config
.
jitter_noise
)
config
[
'num_experts'
],
jitter_noise
=
config
[
'jitter_noise'
])
moe_layer
=
moe
.
MoeLayer
(
experts
,
router
,
train_capacity_factor
=
config
.
train_capacity_factor
,
eval_capacity_factor
=
config
.
eval_capacity_factor
,
max_group_size
=
config
.
max_group_size
,
min_expert_capacity
=
config
.
min_expert_capacity
)
layer
=
moe
.
MoeLayerWithBackbone
(
moe_layer
,
config
.
backbone_d_ff
)
train_capacity_factor
=
config
[
'train_capacity_factor'
]
,
eval_capacity_factor
=
config
[
'eval_capacity_factor'
]
,
max_group_size
=
config
[
'max_group_size'
]
,
min_expert_capacity
=
config
[
'min_expert_capacity'
]
)
layer
=
moe
.
MoeLayerWithBackbone
(
moe_layer
,
config
[
'backbone_d_ff'
]
)
inputs
=
make_input_ones
()
outputs
=
layer
(
inputs
)
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录