namespace Tensorflow.Keras.Optimizers

{

public class AdamW : Adam

{

string name;

float weight_decay;

DeviceDType deType;

List<string> no_decay_params = null;

public AdamW(float learning_rate= 0.001f,

float weight_decay= 0.004f,

float beta_1= 0.9f,

float beta_2= 0.999f,

float epsilon= 1e-7f,

bool amsgrad = false,

List<string> no_decay_params = null,

string name= "AdamW") : base(learning_rate, beta_1, beta_2, epsilon, amsgrad)

{

this.name = name;

this.weight_decay = weight_decay;

this.no_decay_params = no_decay_params;

}

protected Operation _decay_weights_op(IVariableV1 var, float learning_rate, Dictionary<DeviceDType, Dictionary<string, Tensor>> apply_state)

{

bool do_decay = _do_use_weight_decay(var.Name);

if (do_decay) return var.assign_add(

-learning_rate * var.AsTensor() * apply_state[deType]["weight_decay"]);

return tf.no_op();

}

protected bool _do_use_weight_decay(string param_name)

{

// Whether to use L2 weight decay for `param_name`.

if (this.weight_decay == 0)

return false;

if (this.no_decay_params != null)

{

foreach (var name in no_decay_params)

{

if (param_name.Contains(name)) return false;

}

}

return true;

}

protected override Operation _resource_apply_dense(IVariableV1 var, Tensor grad, Dictionary<DeviceDType, Dictionary<string, Tensor>> apply_state)

{

var decay = _decay_weights_op(var, _hyper["learning_rate"], apply_state);

tf.control_dependencies(new[] { decay });

return base._resource_apply_dense(var, grad, apply_state);

}

protected override void _prepare_local(DeviceDType device_dtype, Dictionary<DeviceDType, Dictionary<string, Tensor>> apply_state)

{

this.deType = device_dtype;

base._prepare_local(device_dtype, apply_state);

apply_state[device_dtype]["weight_decay"] = tf.constant(

weight_decay, name: "adam_weight_decay_rate");

}

}

}