Skip to content
GitLab
Projects
Groups
Snippets
/
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
Lucas Miranda
deepOF
Commits
63c23405
Commit
63c23405
authored
Apr 22, 2021
by
lucas_miranda
Browse files
Started implementing annealing mode for KL divergence
parent
e7514b68
Changes
1
Hide whitespace changes
Inline
Side-by-side
deepof/model_utils.py
View file @
63c23405
...
...
@@ -415,11 +415,12 @@ class KLDivergenceLayer(tfpl.KLDivergenceAddLoss):
to the final model loss.
"""
def
__init__
(
self
,
iters
,
warm_up_iters
,
*
args
,
**
kwargs
):
def
__init__
(
self
,
iters
,
warm_up_iters
,
annealing_mode
=
"sigmoid"
,
*
args
,
**
kwargs
):
super
(
KLDivergenceLayer
,
self
).
__init__
(
*
args
,
**
kwargs
)
self
.
is_placeholder
=
True
self
.
_iters
=
iters
self
.
_warm_up_iters
=
warm_up_iters
self
.
_annealing_mode
=
annealing_mode
def
get_config
(
self
):
# pragma: no cover
"""Updates Constraint metadata"""
...
...
@@ -428,6 +429,7 @@ class KLDivergenceLayer(tfpl.KLDivergenceAddLoss):
config
.
update
({
"is_placeholder"
:
self
.
is_placeholder
})
config
.
update
({
"_iters"
:
self
.
_iters
})
config
.
update
({
"_warm_up_iters"
:
self
.
_warm_up_iters
})
config
.
update
({
"_annealing_mode"
:
self
.
_annealing_mode
})
return
config
def
call
(
self
,
distribution_a
):
...
...
@@ -435,9 +437,14 @@ class KLDivergenceLayer(tfpl.KLDivergenceAddLoss):
# Define and update KL weight for warmup
if
self
.
_warm_up_iters
>
0
:
kl_weight
=
tf
.
cast
(
K
.
min
([
self
.
_iters
/
self
.
_warm_up_iters
,
1.0
]),
tf
.
float32
)
if
self
.
_annealing_mode
==
"linear"
:
kl_weight
=
tf
.
cast
(
K
.
min
([
self
.
_iters
/
self
.
_warm_up_iters
,
1.0
]),
tf
.
float32
)
elif
self
.
_annealing_mode
==
"sigmoid"
:
kl_weight
=
0
else
:
raise
NotImplementedError
(
"annealing_mode must be one of 'linear' and 'sigmoid'"
)
else
:
kl_weight
=
tf
.
cast
(
1.0
,
tf
.
float32
)
...
...
Write
Preview
Supports
Markdown
0%
Try again
or
attach a new file
.
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment