Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
曾经的那一瞬间
Models
提交
5fea53a7
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,发现更多精彩内容 >>
提交
5fea53a7
编写于
12月 28, 2021
作者:
H
Hongkun Yu
提交者:
A. Unique TensorFlower
12月 28, 2021
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
Internal change
PiperOrigin-RevId: 418724903
上级
77d9fd62
变更
2
隐藏空白更改
内联
并排
Showing
2 changed file
with
160 addition
and
0 deletion
+160
-0
official/modeling/tf_utils.py
official/modeling/tf_utils.py
+71
-0
official/modeling/tf_utils_test.py
official/modeling/tf_utils_test.py
+89
-0
未找到文件。
official/modeling/tf_utils.py
浏览文件 @
5fea53a7
...
...
@@ -201,3 +201,74 @@ def safe_mean(losses):
total
=
tf
.
reduce_sum
(
losses
)
num_elements
=
tf
.
cast
(
tf
.
size
(
losses
),
dtype
=
losses
.
dtype
)
return
tf
.
math
.
divide_no_nan
(
total
,
num_elements
)
def
get_replica_id
():
"""Gets replica id depending on the environment."""
context
=
tf
.
distribute
.
get_replica_context
()
if
context
is
not
None
:
return
context
.
replica_id_in_sync_group
else
:
raise
RuntimeError
(
"Unknown replica context. The `get_replica_id` method "
"relies on TF 2.x tf.distribute API."
)
def
cross_replica_concat
(
value
,
axis
,
name
=
"cross_replica_concat"
):
"""Concatenates the given `value` across (GPU/TPU) cores, along `axis`.
In general, each core ("replica") will pass a
replica-specific value as `value` (corresponding to some element of a
data-parallel computation taking place across replicas).
The resulting concatenated `Tensor` will have the same shape as `value` for
all dimensions except `axis`, where it will be larger by a factor of the
number of replicas. It will also have the same `dtype` as `value`.
The position of a given replica's `value` within the resulting concatenation
is determined by that replica's replica ID. For
example:
With `value` for replica 0 given as
0 0 0
0 0 0
and `value` for replica 1 given as
1 1 1
1 1 1
the resulting concatenation along axis 0 will be
0 0 0
0 0 0
1 1 1
1 1 1
and this result will be identical across all replicas.
Note that this API only works in TF2 with `tf.distribute`.
Args:
value: The `Tensor` to concatenate across replicas. Each replica will have a
different value for this `Tensor`, and these replica-specific values will
be concatenated.
axis: The axis along which to perform the concatenation as a Python integer
(not a `Tensor`). E.g., `axis=0` to concatenate along the batch dimension.
name: A name for the operation (used to create a name scope).
Returns:
The result of concatenating `value` along `axis` across replicas.
Raises:
RuntimeError: when the batch (0-th) dimension is None.
"""
with
tf
.
name_scope
(
name
):
context
=
tf
.
distribute
.
get_replica_context
()
# Typically this could be hit only if the tensor is derived from a
# dataset with finite epochs and drop_remainder=False, where the last
# batch could of different batch size and then the dim-0 is of dynamic
# shape.
if
value
.
shape
.
as_list
()[
0
]
is
None
:
raise
RuntimeError
(
f
"
{
value
}
has unknown batch."
)
return
context
.
all_gather
(
value
,
axis
=
axis
)
official/modeling/tf_utils_test.py
0 → 100644
浏览文件 @
5fea53a7
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for tf_utils."""
from
absl.testing
import
parameterized
import
numpy
as
np
import
tensorflow
as
tf
from
tensorflow.python.distribute
import
combinations
from
tensorflow.python.distribute
import
strategy_combinations
from
official.modeling
import
tf_utils
def
all_strategy_combinations
():
return
combinations
.
combine
(
strategy
=
[
strategy_combinations
.
cloud_tpu_strategy
,
strategy_combinations
.
mirrored_strategy_with_two_gpus
,
],
mode
=
'eager'
,
)
class
TFUtilsTest
(
tf
.
test
.
TestCase
,
parameterized
.
TestCase
):
@
combinations
.
generate
(
all_strategy_combinations
())
def
test_cross_replica_concat
(
self
,
strategy
):
num_cores
=
strategy
.
num_replicas_in_sync
shape
=
(
2
,
3
,
4
)
def
concat
(
axis
):
@
tf
.
function
def
function
():
replica_value
=
tf
.
fill
(
shape
,
tf_utils
.
get_replica_id
())
return
tf_utils
.
cross_replica_concat
(
replica_value
,
axis
=
axis
)
return
function
def
expected
(
axis
):
values
=
[
np
.
full
(
shape
,
i
)
for
i
in
range
(
num_cores
)]
return
np
.
concatenate
(
values
,
axis
=
axis
)
per_replica_results
=
strategy
.
run
(
concat
(
axis
=
0
))
replica_0_result
=
per_replica_results
.
values
[
0
].
numpy
()
for
value
in
per_replica_results
.
values
[
1
:]:
self
.
assertAllClose
(
value
.
numpy
(),
replica_0_result
)
self
.
assertAllClose
(
replica_0_result
,
expected
(
axis
=
0
))
replica_0_result
=
strategy
.
run
(
concat
(
axis
=
1
)).
values
[
0
].
numpy
()
self
.
assertAllClose
(
replica_0_result
,
expected
(
axis
=
1
))
replica_0_result
=
strategy
.
run
(
concat
(
axis
=
2
)).
values
[
0
].
numpy
()
self
.
assertAllClose
(
replica_0_result
,
expected
(
axis
=
2
))
@
combinations
.
generate
(
all_strategy_combinations
())
def
test_cross_replica_concat_gradient
(
self
,
strategy
):
num_cores
=
strategy
.
num_replicas_in_sync
shape
=
(
10
,
5
)
@
tf
.
function
def
function
():
replica_value
=
tf
.
random
.
normal
(
shape
)
with
tf
.
GradientTape
()
as
tape
:
tape
.
watch
(
replica_value
)
concat_value
=
tf_utils
.
cross_replica_concat
(
replica_value
,
axis
=
0
)
output
=
tf
.
reduce_sum
(
concat_value
)
return
tape
.
gradient
(
output
,
replica_value
)
per_replica_gradients
=
strategy
.
run
(
function
)
for
gradient
in
per_replica_gradients
.
values
:
self
.
assertAllClose
(
gradient
,
num_cores
*
tf
.
ones
(
shape
))
if
__name__
==
'__main__'
:
tf
.
test
.
main
()
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录