From cba4741d1036356971846922347edcf24fccd4d3 Mon Sep 17 00:00:00 2001
From: D-X-Y <280835372@qq.com>
Date: Sun, 5 Jul 2020 23:19:25 +0000
Subject: [PATCH] Remove TF codes
---
exps/experimental/test-ww-bench.py | 2 +-
lib/tf_models/__init__.py | 32 --
lib/tf_models/cell_operations.py | 150 -------
lib/tf_models/cell_searchs/__init__.py | 8 -
lib/tf_models/cell_searchs/search_cells.py | 50 ---
.../cell_searchs/search_model_darts.py | 83 ----
.../cell_searchs/search_model_gdas.py | 99 ----
lib/tf_optimizers/__init__.py | 1 -
lib/tf_optimizers/weight_decay_optimizers.py | 422 ------------------
9 files changed, 1 insertion(+), 846 deletions(-)
delete mode 100644 lib/tf_models/__init__.py
delete mode 100644 lib/tf_models/cell_operations.py
delete mode 100644 lib/tf_models/cell_searchs/__init__.py
delete mode 100644 lib/tf_models/cell_searchs/search_cells.py
delete mode 100644 lib/tf_models/cell_searchs/search_model_darts.py
delete mode 100644 lib/tf_models/cell_searchs/search_model_gdas.py
delete mode 100644 lib/tf_optimizers/__init__.py
delete mode 100644 lib/tf_optimizers/weight_decay_optimizers.py
diff --git a/exps/experimental/test-ww-bench.py b/exps/experimental/test-ww-bench.py
index 4351571..a07dfa0 100644
--- a/exps/experimental/test-ww-bench.py
+++ b/exps/experimental/test-ww-bench.py
@@ -70,7 +70,7 @@ def evaluate(api, weight_dir, data: str):
ok += 1
norms.append(cur_norm)
# query the accuracy
- info = meta_info.get_metrics(data, 'ori-test', iepoch=None, is_random=777)
+ info = meta_info.get_metrics(data, 'ori-test', iepoch=None, is_random=888 if isinstance(api, NASBench201API) else 777)
accuracies.append(info['accuracy'])
del net, meta_info
# print the information
diff --git a/lib/tf_models/__init__.py b/lib/tf_models/__init__.py
deleted file mode 100644
index d402913..0000000
--- a/lib/tf_models/__init__.py
+++ /dev/null
@@ -1,32 +0,0 @@
-##################################################
-# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
-##################################################
-import torch
-from os import path as osp
-
-__all__ = ['get_cell_based_tiny_net', 'get_search_spaces']
-
-
-# the cell-based NAS models
-def get_cell_based_tiny_net(config):
- group_names = ['GDAS', 'DARTS']
- if config.name in group_names:
- from .cell_searchs import nas_super_nets
- from .cell_operations import SearchSpaceNames
- if isinstance(config.space, str): search_space = SearchSpaceNames[config.space]
- else: search_space = config.space
- return nas_super_nets[config.name](
- config.C, config.N, config.max_nodes,
- config.num_classes, search_space, config.affine)
- else:
- raise ValueError('invalid network name : {:}'.format(config.name))
-
-
-# obtain the search space, i.e., a dict mapping the operation name into a python-function for this op
-def get_search_spaces(xtype, name):
- if xtype == 'cell':
- from .cell_operations import SearchSpaceNames
- assert name in SearchSpaceNames, 'invalid name [{:}] in {:}'.format(name, SearchSpaceNames.keys())
- return SearchSpaceNames[name]
- else:
- raise ValueError('invalid search-space type is {:}'.format(xtype))
diff --git a/lib/tf_models/cell_operations.py b/lib/tf_models/cell_operations.py
deleted file mode 100644
index 78f1c2a..0000000
--- a/lib/tf_models/cell_operations.py
+++ /dev/null
@@ -1,150 +0,0 @@
-##################################################
-# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
-##################################################
-import tensorflow as tf
-
-__all__ = ['OPS', 'ResNetBasicblock', 'SearchSpaceNames']
-
-OPS = {
- 'none' : lambda C_in, C_out, stride, affine: Zero(C_in, C_out, stride),
- 'avg_pool_3x3': lambda C_in, C_out, stride, affine: POOLING(C_in, C_out, stride, 'avg', affine),
- 'nor_conv_1x1': lambda C_in, C_out, stride, affine: ReLUConvBN(C_in, C_out, 1, stride, affine),
- 'nor_conv_3x3': lambda C_in, C_out, stride, affine: ReLUConvBN(C_in, C_out, 3, stride, affine),
- 'nor_conv_5x5': lambda C_in, C_out, stride, affine: ReLUConvBN(C_in, C_out, 5, stride, affine),
- 'skip_connect': lambda C_in, C_out, stride, affine: Identity(C_in, C_out, stride) if stride == 1 and C_in == C_out else FactorizedReduce(C_in, C_out, stride, affine)
-}
-
-NAS_BENCH_201 = ['none', 'skip_connect', 'nor_conv_1x1', 'nor_conv_3x3', 'avg_pool_3x3']
-
-SearchSpaceNames = {
- 'nas-bench-201': NAS_BENCH_201,
- }
-
-
-class POOLING(tf.keras.layers.Layer):
-
- def __init__(self, C_in, C_out, stride, mode, affine):
- super(POOLING, self).__init__()
- if C_in == C_out:
- self.preprocess = None
- else:
- self.preprocess = ReLUConvBN(C_in, C_out, 1, 1, affine)
- if mode == 'avg' : self.op = tf.keras.layers.AvgPool2D((3,3), strides=stride, padding='same')
- elif mode == 'max': self.op = tf.keras.layers.MaxPool2D((3,3), strides=stride, padding='same')
- else : raise ValueError('Invalid mode={:} in POOLING'.format(mode))
-
- def call(self, inputs, training):
- if self.preprocess: x = self.preprocess(inputs)
- else : x = inputs
- return self.op(x)
-
-
-class Identity(tf.keras.layers.Layer):
- def __init__(self, C_in, C_out, stride):
- super(Identity, self).__init__()
- if C_in != C_out or stride != 1:
- self.layer = tf.keras.layers.Conv2D(C_out, 3, stride, padding='same', use_bias=False)
- else:
- self.layer = None
-
- def call(self, inputs, training):
- x = inputs
- if self.layer is not None:
- x = self.layer(x)
- return x
-
-
-
-class Zero(tf.keras.layers.Layer):
- def __init__(self, C_in, C_out, stride):
- super(Zero, self).__init__()
- if C_in != C_out:
- self.layer = tf.keras.layers.Conv2D(C_out, 1, stride, padding='same', use_bias=False)
- elif stride != 1:
- self.layer = tf.keras.layers.AvgPool2D((stride,stride), None, padding="same")
- else:
- self.layer = None
-
- def call(self, inputs, training):
- x = tf.zeros_like(inputs)
- if self.layer is not None:
- x = self.layer(x)
- return x
-
-
-class ReLUConvBN(tf.keras.layers.Layer):
- def __init__(self, C_in, C_out, kernel_size, strides, affine):
- super(ReLUConvBN, self).__init__()
- self.C_in = C_in
- self.relu = tf.keras.activations.relu
- self.conv = tf.keras.layers.Conv2D(C_out, kernel_size, strides, padding='same', use_bias=False)
- self.bn = tf.keras.layers.BatchNormalization(center=affine, scale=affine)
-
- def call(self, inputs, training):
- x = self.relu(inputs)
- x = self.conv(x)
- x = self.bn(x, training)
- return x
-
-
-class FactorizedReduce(tf.keras.layers.Layer):
- def __init__(self, C_in, C_out, stride, affine):
- assert output_filters % 2 == 0, ('Need even number of filters when using this factorized reduction.')
- self.stride == stride
- self.relu = tf.keras.activations.relu
- if stride == 1:
- self.layer = tf.keras.Sequential([
- tf.keras.layers.Conv2D(C_out, 1, strides, padding='same', use_bias=False),
- tf.keras.layers.BatchNormalization(center=affine, scale=affine)])
- elif stride == 2:
- stride_spec = [1, stride, stride, 1] # data_format == 'NHWC'
- self.layer1 = tf.keras.layers.Conv2D(C_out//2, 1, strides, padding='same', use_bias=False)
- self.layer2 = tf.keras.layers.Conv2D(C_out//2, 1, strides, padding='same', use_bias=False)
- self.bn = tf.keras.layers.BatchNormalization(center=affine, scale=affine)
- else:
- raise ValueError('invalid stride={:}'.format(stride))
-
- def call(self, inputs, training):
- x = self.relu(inputs)
- if self.stride == 1:
- return self.layer(x, training)
- else:
- path1 = x
- path2 = tf.pad(x, [[0, 0], [0, 1], [0, 1], [0, 0]])[:, 1:, 1:, :] # data_format == 'NHWC'
- x1 = self.layer1(path1)
- x2 = self.layer2(path2)
- final_path = tf.concat(values=[x1, x2], axis=3)
- return self.bn(final_path)
-
-
-class ResNetBasicblock(tf.keras.layers.Layer):
-
- def __init__(self, inplanes, planes, stride, affine=True):
- super(ResNetBasicblock, self).__init__()
- assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
- self.conv_a = ReLUConvBN(inplanes, planes, 3, stride, affine)
- self.conv_b = ReLUConvBN( planes, planes, 3, 1, affine)
- if stride == 2:
- self.downsample = tf.keras.Sequential([
- tf.keras.layers.AvgPool2D((stride,stride), None, padding="same"),
- tf.keras.layers.Conv2D(planes, 1, 1, padding='same', use_bias=False)])
- elif inplanes != planes:
- self.downsample = ReLUConvBN(inplanes, planes, 1, stride, affine)
- else:
- self.downsample = None
- self.addition = tf.keras.layers.Add()
- self.in_dim = inplanes
- self.out_dim = planes
- self.stride = stride
- self.num_conv = 2
-
- def call(self, inputs, training):
-
- basicblock = self.conv_a(inputs, training)
- basicblock = self.conv_b(basicblock, training)
-
- if self.downsample is not None:
- residual = self.downsample(inputs)
- else:
- residual = inputs
- return self.addition([residual, basicblock])
diff --git a/lib/tf_models/cell_searchs/__init__.py b/lib/tf_models/cell_searchs/__init__.py
deleted file mode 100644
index 717fbe4..0000000
--- a/lib/tf_models/cell_searchs/__init__.py
+++ /dev/null
@@ -1,8 +0,0 @@
-##################################################
-# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
-##################################################
-from .search_model_gdas import TinyNetworkGDAS
-from .search_model_darts import TinyNetworkDARTS
-
-nas_super_nets = {'GDAS' : TinyNetworkGDAS,
- 'DARTS': TinyNetworkDARTS}
diff --git a/lib/tf_models/cell_searchs/search_cells.py b/lib/tf_models/cell_searchs/search_cells.py
deleted file mode 100644
index 83de6b3..0000000
--- a/lib/tf_models/cell_searchs/search_cells.py
+++ /dev/null
@@ -1,50 +0,0 @@
-##################################################
-# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
-##################################################
-import math, random
-import tensorflow as tf
-from copy import deepcopy
-from ..cell_operations import OPS
-
-
-class NAS201SearchCell(tf.keras.layers.Layer):
-
- def __init__(self, C_in, C_out, stride, max_nodes, op_names, affine=False):
- super(NAS201SearchCell, self).__init__()
-
- self.op_names = deepcopy(op_names)
- self.max_nodes = max_nodes
- self.in_dim = C_in
- self.out_dim = C_out
- self.edge_keys = []
- for i in range(1, max_nodes):
- for j in range(i):
- node_str = '{:}<-{:}'.format(i, j)
- if j == 0:
- xlists = [OPS[op_name](C_in , C_out, stride, affine) for op_name in op_names]
- else:
- xlists = [OPS[op_name](C_in , C_out, 1, affine) for op_name in op_names]
- for k, op in enumerate(xlists):
- setattr(self, '{:}.{:}'.format(node_str, k), op)
- self.edge_keys.append( node_str )
- self.edge_keys = sorted(self.edge_keys)
- self.edge2index = {key:i for i, key in enumerate(self.edge_keys)}
- self.num_edges = len(self.edge_keys)
-
- def call(self, inputs, weightss, training):
- w_lst = tf.split(weightss, self.num_edges, 0)
- nodes = [inputs]
- for i in range(1, self.max_nodes):
- inter_nodes = []
- for j in range(i):
- node_str = '{:}<-{:}'.format(i, j)
- edge_idx = self.edge2index[node_str]
- op_outps = []
- for k, op_name in enumerate(self.op_names):
- op = getattr(self, '{:}.{:}'.format(node_str, k))
- op_outps.append( op(nodes[j], training) )
- stack_op_outs = tf.stack(op_outps, axis=-1)
- weighted_sums = tf.math.multiply(stack_op_outs, w_lst[edge_idx])
- inter_nodes.append( tf.math.reduce_sum(weighted_sums, axis=-1) )
- nodes.append( tf.math.add_n(inter_nodes) )
- return nodes[-1]
diff --git a/lib/tf_models/cell_searchs/search_model_darts.py b/lib/tf_models/cell_searchs/search_model_darts.py
deleted file mode 100644
index ad05b8b..0000000
--- a/lib/tf_models/cell_searchs/search_model_darts.py
+++ /dev/null
@@ -1,83 +0,0 @@
-import tensorflow as tf
-import numpy as np
-from copy import deepcopy
-from ..cell_operations import ResNetBasicblock
-from .search_cells import NAS201SearchCell as SearchCell
-
-
-class TinyNetworkDARTS(tf.keras.Model):
-
- def __init__(self, C, N, max_nodes, num_classes, search_space, affine):
- super(TinyNetworkDARTS, self).__init__()
- self._C = C
- self._layerN = N
- self.max_nodes = max_nodes
- self.stem = tf.keras.Sequential([
- tf.keras.layers.Conv2D(16, 3, 1, padding='same', use_bias=False),
- tf.keras.layers.BatchNormalization()], name='stem')
-
- layer_channels = [C ] * N + [C*2 ] + [C*2 ] * N + [C*4 ] + [C*4 ] * N
- layer_reductions = [False] * N + [True] + [False] * N + [True] + [False] * N
-
- C_prev, num_edge, edge2index = C, None, None
- for index, (C_curr, reduction) in enumerate(zip(layer_channels, layer_reductions)):
- cell_prefix = 'cell-{:03d}'.format(index)
- #with tf.name_scope(cell_prefix) as scope:
- if reduction:
- cell = ResNetBasicblock(C_prev, C_curr, 2)
- else:
- cell = SearchCell(C_prev, C_curr, 1, max_nodes, search_space, affine)
- if num_edge is None: num_edge, edge2index = cell.num_edges, cell.edge2index
- else: assert num_edge == cell.num_edges and edge2index == cell.edge2index, 'invalid {:} vs. {:}.'.format(num_edge, cell.num_edges)
- C_prev = cell.out_dim
- setattr(self, cell_prefix, cell)
- self.num_layers = len(layer_reductions)
- self.op_names = deepcopy( search_space )
- self.edge2index = edge2index
- self.num_edge = num_edge
- self.lastact = tf.keras.Sequential([
- tf.keras.layers.BatchNormalization(),
- tf.keras.layers.ReLU(),
- tf.keras.layers.GlobalAvgPool2D(),
- tf.keras.layers.Flatten(),
- tf.keras.layers.Dense(num_classes, activation='softmax')], name='lastact')
- #self.arch_parameters = nn.Parameter( 1e-3*torch.randn(num_edge, len(search_space)) )
- arch_init = tf.random_normal_initializer(mean=0, stddev=0.001)
- self.arch_parameters = tf.Variable(initial_value=arch_init(shape=(num_edge, len(search_space)), dtype='float32'), trainable=True, name='arch-encoding')
-
- def get_alphas(self):
- xlist = self.trainable_variables
- return [x for x in xlist if 'arch-encoding' in x.name]
-
- def get_weights(self):
- xlist = self.trainable_variables
- return [x for x in xlist if 'arch-encoding' not in x.name]
-
- def get_np_alphas(self):
- arch_nps = self.arch_parameters.numpy()
- arch_ops = np.exp(arch_nps) / np.sum(np.exp(arch_nps), axis=-1, keepdims=True)
- return arch_ops
-
- def genotype(self):
- genotypes, arch_nps = [], self.arch_parameters.numpy()
- for i in range(1, self.max_nodes):
- xlist = []
- for j in range(i):
- node_str = '{:}<-{:}'.format(i, j)
- weights = arch_nps[ self.edge2index[node_str] ]
- op_name = self.op_names[ weights.argmax().item() ]
- xlist.append((op_name, j))
- genotypes.append( tuple(xlist) )
- return genotypes
-
- def call(self, inputs, training):
- weightss = tf.nn.softmax(self.arch_parameters, axis=1)
- feature = self.stem(inputs, training)
- for idx in range(self.num_layers):
- cell = getattr(self, 'cell-{:03d}'.format(idx))
- if isinstance(cell, SearchCell):
- feature = cell.call(feature, weightss, training)
- else:
- feature = cell(feature, training)
- logits = self.lastact(feature, training)
- return logits
diff --git a/lib/tf_models/cell_searchs/search_model_gdas.py b/lib/tf_models/cell_searchs/search_model_gdas.py
deleted file mode 100644
index d10bb19..0000000
--- a/lib/tf_models/cell_searchs/search_model_gdas.py
+++ /dev/null
@@ -1,99 +0,0 @@
-###########################################################################
-# Searching for A Robust Neural Architecture in Four GPU Hours, CVPR 2019 #
-###########################################################################
-import tensorflow as tf
-import numpy as np
-from copy import deepcopy
-from ..cell_operations import ResNetBasicblock
-from .search_cells import NAS201SearchCell as SearchCell
-
-
-def sample_gumbel(shape, eps=1e-20):
- U = tf.random.uniform(shape, minval=0, maxval=1)
- return -tf.math.log(-tf.math.log(U + eps) + eps)
-
-
-def gumbel_softmax(logits, temperature):
- gumbel_softmax_sample = logits + sample_gumbel(tf.shape(logits))
- y = tf.nn.softmax(gumbel_softmax_sample / temperature)
- return y
-
-
-class TinyNetworkGDAS(tf.keras.Model):
-
- def __init__(self, C, N, max_nodes, num_classes, search_space, affine):
- super(TinyNetworkGDAS, self).__init__()
- self._C = C
- self._layerN = N
- self.max_nodes = max_nodes
- self.stem = tf.keras.Sequential([
- tf.keras.layers.Conv2D(16, 3, 1, padding='same', use_bias=False),
- tf.keras.layers.BatchNormalization()], name='stem')
-
- layer_channels = [C ] * N + [C*2 ] + [C*2 ] * N + [C*4 ] + [C*4 ] * N
- layer_reductions = [False] * N + [True] + [False] * N + [True] + [False] * N
-
- C_prev, num_edge, edge2index = C, None, None
- for index, (C_curr, reduction) in enumerate(zip(layer_channels, layer_reductions)):
- cell_prefix = 'cell-{:03d}'.format(index)
- #with tf.name_scope(cell_prefix) as scope:
- if reduction:
- cell = ResNetBasicblock(C_prev, C_curr, 2)
- else:
- cell = SearchCell(C_prev, C_curr, 1, max_nodes, search_space, affine)
- if num_edge is None: num_edge, edge2index = cell.num_edges, cell.edge2index
- else: assert num_edge == cell.num_edges and edge2index == cell.edge2index, 'invalid {:} vs. {:}.'.format(num_edge, cell.num_edges)
- C_prev = cell.out_dim
- setattr(self, cell_prefix, cell)
- self.num_layers = len(layer_reductions)
- self.op_names = deepcopy( search_space )
- self.edge2index = edge2index
- self.num_edge = num_edge
- self.lastact = tf.keras.Sequential([
- tf.keras.layers.BatchNormalization(),
- tf.keras.layers.ReLU(),
- tf.keras.layers.GlobalAvgPool2D(),
- tf.keras.layers.Flatten(),
- tf.keras.layers.Dense(num_classes, activation='softmax')], name='lastact')
- #self.arch_parameters = nn.Parameter( 1e-3*torch.randn(num_edge, len(search_space)) )
- arch_init = tf.random_normal_initializer(mean=0, stddev=0.001)
- self.arch_parameters = tf.Variable(initial_value=arch_init(shape=(num_edge, len(search_space)), dtype='float32'), trainable=True, name='arch-encoding')
-
- def get_alphas(self):
- xlist = self.trainable_variables
- return [x for x in xlist if 'arch-encoding' in x.name]
-
- def get_weights(self):
- xlist = self.trainable_variables
- return [x for x in xlist if 'arch-encoding' not in x.name]
-
- def get_np_alphas(self):
- arch_nps = self.arch_parameters.numpy()
- arch_ops = np.exp(arch_nps) / np.sum(np.exp(arch_nps), axis=-1, keepdims=True)
- return arch_ops
-
- def genotype(self):
- genotypes, arch_nps = [], self.arch_parameters.numpy()
- for i in range(1, self.max_nodes):
- xlist = []
- for j in range(i):
- node_str = '{:}<-{:}'.format(i, j)
- weights = arch_nps[ self.edge2index[node_str] ]
- op_name = self.op_names[ weights.argmax().item() ]
- xlist.append((op_name, j))
- genotypes.append( tuple(xlist) )
- return genotypes
-
- #
- def call(self, inputs, tau, training):
- weightss = tf.cond(tau < 0, lambda: tf.nn.softmax(self.arch_parameters, axis=1),
- lambda: gumbel_softmax(tf.math.log_softmax(self.arch_parameters, axis=1), tau))
- feature = self.stem(inputs, training)
- for idx in range(self.num_layers):
- cell = getattr(self, 'cell-{:03d}'.format(idx))
- if isinstance(cell, SearchCell):
- feature = cell.call(feature, weightss, training)
- else:
- feature = cell(feature, training)
- logits = self.lastact(feature, training)
- return logits
diff --git a/lib/tf_optimizers/__init__.py b/lib/tf_optimizers/__init__.py
deleted file mode 100644
index c72fe17..0000000
--- a/lib/tf_optimizers/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from .weight_decay_optimizers import AdamW, SGDW
diff --git a/lib/tf_optimizers/weight_decay_optimizers.py b/lib/tf_optimizers/weight_decay_optimizers.py
deleted file mode 100644
index b4e72dc..0000000
--- a/lib/tf_optimizers/weight_decay_optimizers.py
+++ /dev/null
@@ -1,422 +0,0 @@
-# Copyright 2019 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.
-# ==============================================================================
-"""Base class to make optimizers weight decay ready."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import tensorflow as tf
-
-
-class DecoupledWeightDecayExtension(object):
- """This class allows to extend optimizers with decoupled weight decay.
-
- It implements the decoupled weight decay described by Loshchilov & Hutter
- (https://arxiv.org/pdf/1711.05101.pdf), in which the weight decay is
- decoupled from the optimization steps w.r.t. to the loss function.
- For SGD variants, this simplifies hyperparameter search since it decouples
- the settings of weight decay and learning rate.
- For adaptive gradient algorithms, it regularizes variables with large
- gradients more than L2 regularization would, which was shown to yield
- better training loss and generalization error in the paper above.
-
- This class alone is not an optimizer but rather extends existing
- optimizers with decoupled weight decay. We explicitly define the two
- examples used in the above paper (SGDW and AdamW), but in general this
- can extend any OptimizerX by using
- `extend_with_decoupled_weight_decay(
- OptimizerX, weight_decay=weight_decay)`.
- In order for it to work, it must be the first class the Optimizer with
- weight decay inherits from, e.g.
-
- ```python
- class AdamW(DecoupledWeightDecayExtension, tf.keras.optimizers.Adam):
- def __init__(self, weight_decay, *args, **kwargs):
- super(AdamW, self).__init__(weight_decay, *args, **kwargs).
- ```
-
- Note: this extension decays weights BEFORE applying the update based
- on the gradient, i.e. this extension only has the desired behaviour for
- optimizers which do not depend on the value of'var' in the update step!
-
- Note: when applying a decay to the learning rate, be sure to manually apply
- the decay to the `weight_decay` as well. For example:
-
- ```python
- step = tf.Variable(0, trainable=False)
- schedule = tf.optimizers.schedules.PiecewiseConstantDecay(
- [10000, 15000], [1e-0, 1e-1, 1e-2])
- # lr and wd can be a function or a tensor
- lr = 1e-1 * schedule(step)
- wd = lambda: 1e-4 * schedule(step)
-
- # ...
-
- optimizer = tfa.optimizers.AdamW(learning_rate=lr, weight_decay=wd)
- ```
- """
-
- def __init__(self, weight_decay, **kwargs):
- """Extension class that adds weight decay to an optimizer.
-
- Args:
- weight_decay: A `Tensor` or a floating point value, the factor by
- which a variable is decayed in the update step.
- **kwargs: Optional list or tuple or set of `Variable` objects to
- decay.
- """
- wd = kwargs.pop('weight_decay', weight_decay)
- super(DecoupledWeightDecayExtension, self).__init__(**kwargs)
- self._decay_var_list = None # is set in minimize or apply_gradients
- self._set_hyper('weight_decay', wd)
-
- def get_config(self):
- config = super(DecoupledWeightDecayExtension, self).get_config()
- config.update({
- 'weight_decay':
- self._serialize_hyperparameter('weight_decay'),
- })
- return config
-
- def minimize(self,
- loss,
- var_list,
- grad_loss=None,
- name=None,
- decay_var_list=None):
- """Minimize `loss` by updating `var_list`.
-
- This method simply computes gradient using `tf.GradientTape` and calls
- `apply_gradients()`. If you want to process the gradient before
- applying then call `tf.GradientTape` and `apply_gradients()` explicitly
- instead of using this function.
-
- Args:
- loss: A callable taking no arguments which returns the value to
- minimize.
- var_list: list or tuple of `Variable` objects to update to
- minimize `loss`, or a callable returning the list or tuple of
- `Variable` objects. Use callable when the variable list would
- otherwise be incomplete before `minimize` since the variables
- are created at the first time `loss` is called.
- grad_loss: Optional. A `Tensor` holding the gradient computed for
- `loss`.
- decay_var_list: Optional list of variables to be decayed. Defaults
- to all variables in var_list.
- name: Optional name for the returned operation.
- Returns:
- An Operation that updates the variables in `var_list`. If
- `global_step` was not `None`, that operation also increments
- `global_step`.
- Raises:
- ValueError: If some of the variables are not `Variable` objects.
- """
- self._decay_var_list = set(decay_var_list) if decay_var_list else False
- return super(DecoupledWeightDecayExtension, self).minimize(
- loss, var_list=var_list, grad_loss=grad_loss, name=name)
-
- def apply_gradients(self, grads_and_vars, name=None, decay_var_list=None):
- """Apply gradients to variables.
-
- This is the second part of `minimize()`. It returns an `Operation` that
- applies gradients.
-
- Args:
- grads_and_vars: List of (gradient, variable) pairs.
- name: Optional name for the returned operation. Default to the
- name passed to the `Optimizer` constructor.
- decay_var_list: Optional list of variables to be decayed. Defaults
- to all variables in var_list.
- Returns:
- An `Operation` that applies the specified gradients. If
- `global_step` was not None, that operation also increments
- `global_step`.
- Raises:
- TypeError: If `grads_and_vars` is malformed.
- ValueError: If none of the variables have gradients.
- """
- self._decay_var_list = set(decay_var_list) if decay_var_list else False
- return super(DecoupledWeightDecayExtension, self).apply_gradients(
- grads_and_vars, name=name)
-
- def _decay_weights_op(self, var):
- if not self._decay_var_list or var in self._decay_var_list:
- return var.assign_sub(
- self._get_hyper('weight_decay', var.dtype) * var,
- self._use_locking)
- return tf.no_op()
-
- def _decay_weights_sparse_op(self, var, indices):
- if not self._decay_var_list or var in self._decay_var_list:
- update = (-self._get_hyper('weight_decay', var.dtype) * tf.gather(
- var, indices))
- return self._resource_scatter_add(var, indices, update)
- return tf.no_op()
-
- # Here, we overwrite the apply functions that the base optimizer calls.
- # super().apply_x resolves to the apply_x function of the BaseOptimizer.
-
- def _resource_apply_dense(self, grad, var):
- with tf.control_dependencies([self._decay_weights_op(var)]):
- return super(DecoupledWeightDecayExtension,
- self)._resource_apply_dense(grad, var)
-
- def _resource_apply_sparse(self, grad, var, indices):
- decay_op = self._decay_weights_sparse_op(var, indices)
- with tf.control_dependencies([decay_op]):
- return super(DecoupledWeightDecayExtension,
- self)._resource_apply_sparse(grad, var, indices)
-
-
-def extend_with_decoupled_weight_decay(base_optimizer):
- """Factory function returning an optimizer class with decoupled weight
- decay.
-
- Returns an optimizer class. An instance of the returned class computes the
- update step of `base_optimizer` and additionally decays the weights.
- E.g., the class returned by
- `extend_with_decoupled_weight_decay(tf.keras.optimizers.Adam)` is
- equivalent to `tfa.optimizers.AdamW`.
-
- The API of the new optimizer class slightly differs from the API of the
- base optimizer:
- - The first argument to the constructor is the weight decay rate.
- - `minimize` and `apply_gradients` accept the optional keyword argument
- `decay_var_list`, which specifies the variables that should be decayed.
- If `None`, all variables that are optimized are decayed.
-
- Usage example:
- ```python
- # MyAdamW is a new class
- MyAdamW = extend_with_decoupled_weight_decay(tf.keras.optimizers.Adam)
- # Create a MyAdamW object
- optimizer = MyAdamW(weight_decay=0.001, learning_rate=0.001)
- # update var1, var2 but only decay var1
- optimizer.minimize(loss, var_list=[var1, var2], decay_variables=[var1])
-
- Note: this extension decays weights BEFORE applying the update based
- on the gradient, i.e. this extension only has the desired behaviour for
- optimizers which do not depend on the value of 'var' in the update step!
-
- Note: when applying a decay to the learning rate, be sure to manually apply
- the decay to the `weight_decay` as well. For example:
-
- ```python
- step = tf.Variable(0, trainable=False)
- schedule = tf.optimizers.schedules.PiecewiseConstantDecay(
- [10000, 15000], [1e-0, 1e-1, 1e-2])
- # lr and wd can be a function or a tensor
- lr = 1e-1 * schedule(step)
- wd = lambda: 1e-4 * schedule(step)
-
- # ...
-
- optimizer = tfa.optimizers.AdamW(learning_rate=lr, weight_decay=wd)
- ```
-
- Note: you might want to register your own custom optimizer using
- `tf.keras.utils.get_custom_objects()`.
-
- Args:
- base_optimizer: An optimizer class that inherits from
- tf.optimizers.Optimizer.
-
- Returns:
- A new optimizer class that inherits from DecoupledWeightDecayExtension
- and base_optimizer.
- """
-
- class OptimizerWithDecoupledWeightDecay(DecoupledWeightDecayExtension,
- base_optimizer):
- """Base_optimizer with decoupled weight decay.
-
- This class computes the update step of `base_optimizer` and
- additionally decays the variable with the weight decay being
- decoupled from the optimization steps w.r.t. to the loss
- function, as described by Loshchilov & Hutter
- (https://arxiv.org/pdf/1711.05101.pdf). For SGD variants, this
- simplifies hyperparameter search since it decouples the settings
- of weight decay and learning rate. For adaptive gradient
- algorithms, it regularizes variables with large gradients more
- than L2 regularization would, which was shown to yield better
- training loss and generalization error in the paper above.
- """
-
- def __init__(self, weight_decay, *args, **kwargs):
- # super delegation is necessary here
- super(OptimizerWithDecoupledWeightDecay, self).__init__(
- weight_decay, *args, **kwargs)
-
- return OptimizerWithDecoupledWeightDecay
-
-
-class SGDW(DecoupledWeightDecayExtension, tf.keras.optimizers.SGD):
- """Optimizer that implements the Momentum algorithm with weight_decay.
-
- This is an implementation of the SGDW optimizer described in "Decoupled
- Weight Decay Regularization" by Loshchilov & Hutter
- (https://arxiv.org/abs/1711.05101)
- ([pdf])(https://arxiv.org/pdf/1711.05101.pdf).
- It computes the update step of `tf.keras.optimizers.SGD` and additionally
- decays the variable. Note that this is different from adding
- L2 regularization on the variables to the loss. Decoupling the weight decay
- from other hyperparameters (in particular the learning rate) simplifies
- hyperparameter search.
-
- For further information see the documentation of the SGD Optimizer.
-
- This optimizer can also be instantiated as
- ```python
- extend_with_decoupled_weight_decay(tf.keras.optimizers.SGD,
- weight_decay=weight_decay)
- ```
-
- Note: when applying a decay to the learning rate, be sure to manually apply
- the decay to the `weight_decay` as well. For example:
-
- ```python
- step = tf.Variable(0, trainable=False)
- schedule = tf.optimizers.schedules.PiecewiseConstantDecay(
- [10000, 15000], [1e-0, 1e-1, 1e-2])
- # lr and wd can be a function or a tensor
- lr = 1e-1 * schedule(step)
- wd = lambda: 1e-4 * schedule(step)
-
- # ...
-
- optimizer = tfa.optimizers.SGDW(
- learning_rate=lr, weight_decay=wd, momentum=0.9)
- ```
- """
-
- def __init__(self,
- weight_decay,
- learning_rate=0.001,
- momentum=0.0,
- nesterov=False,
- name='SGDW',
- **kwargs):
- """Construct a new SGDW optimizer.
-
- For further information see the documentation of the SGD Optimizer.
-
- Args:
- learning_rate: float hyperparameter >= 0. Learning rate.
- momentum: float hyperparameter >= 0 that accelerates SGD in the
- relevant direction and dampens oscillations.
- nesterov: boolean. Whether to apply Nesterov momentum.
- name: Optional name prefix for the operations created when applying
- gradients. Defaults to 'SGD'.
- **kwargs: keyword arguments. Allowed to be {`clipnorm`,
- `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by
- norm; `clipvalue` is clip gradients by value, `decay` is
- included for backward compatibility to allow time inverse decay
- of learning rate. `lr` is included for backward compatibility,
- recommended to use `learning_rate` instead.
- """
- super(SGDW, self).__init__(
- weight_decay,
- learning_rate=learning_rate,
- momentum=momentum,
- nesterov=nesterov,
- name=name,
- **kwargs)
-
-
-class AdamW(DecoupledWeightDecayExtension, tf.keras.optimizers.Adam):
- """Optimizer that implements the Adam algorithm with weight decay.
-
- This is an implementation of the AdamW optimizer described in "Decoupled
- Weight Decay Regularization" by Loshchilov & Hutter
- (https://arxiv.org/abs/1711.05101)
- ([pdf])(https://arxiv.org/pdf/1711.05101.pdf).
-
- It computes the update step of `tf.keras.optimizers.Adam` and additionally
- decays the variable. Note that this is different from adding L2
- regularization on the variables to the loss: it regularizes variables with
- large gradients more than L2 regularization would, which was shown to yield
- better training loss and generalization error in the paper above.
-
- For further information see the documentation of the Adam Optimizer.
-
- This optimizer can also be instantiated as
- ```python
- extend_with_decoupled_weight_decay(tf.keras.optimizers.Adam,
- weight_decay=weight_decay)
- ```
-
- Note: when applying a decay to the learning rate, be sure to manually apply
- the decay to the `weight_decay` as well. For example:
-
- ```python
- step = tf.Variable(0, trainable=False)
- schedule = tf.optimizers.schedules.PiecewiseConstantDecay(
- [10000, 15000], [1e-0, 1e-1, 1e-2])
- # lr and wd can be a function or a tensor
- lr = 1e-1 * schedule(step)
- wd = lambda: 1e-4 * schedule(step)
-
- # ...
-
- optimizer = tfa.optimizers.AdamW(learning_rate=lr, weight_decay=wd)
- ```
- """
-
- def __init__(self,
- weight_decay,
- learning_rate=0.001,
- beta_1=0.9,
- beta_2=0.999,
- epsilon=1e-07,
- amsgrad=False,
- name="AdamW",
- **kwargs):
- """Construct a new AdamW optimizer.
-
- For further information see the documentation of the Adam Optimizer.
-
- Args:
- weight_decay: A Tensor or a floating point value. The weight decay.
- learning_rate: A Tensor or a floating point value. The learning
- rate.
- beta_1: A float value or a constant float tensor. The exponential
- decay rate for the 1st moment estimates.
- beta_2: A float value or a constant float tensor. The exponential
- decay rate for the 2nd moment estimates.
- epsilon: A small constant for numerical stability. This epsilon is
- "epsilon hat" in the Kingma and Ba paper (in the formula just
- before Section 2.1), not the epsilon in Algorithm 1 of the
- paper.
- amsgrad: boolean. Whether to apply AMSGrad variant of this
- algorithm from the paper "On the Convergence of Adam and
- beyond".
- name: Optional name for the operations created when applying
- gradients. Defaults to "AdamW".
- **kwargs: keyword arguments. Allowed to be {`clipnorm`,
- `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by
- norm; `clipvalue` is clip gradients by value, `decay` is
- included for backward compatibility to allow time inverse decay
- of learning rate. `lr` is included for backward compatibility,
- recommended to use `learning_rate` instead.
- """
- super(AdamW, self).__init__(
- weight_decay,
- learning_rate=learning_rate,
- beta_1=beta_1,
- beta_2=beta_2,
- epsilon=epsilon,
- amsgrad=amsgrad,
- name=name,
- **kwargs)