TensorFlow.NET/src/TensorFlowNET.Core/Training/Trackable.cs at master · hartoto2/TensorFlow.NET · GitHub

History

361 lines (328 loc) · 15.4 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

/*****************************************************************************

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.

******************************************************************************/

using OneOf;

using System;

using System.Collections.Generic;

using System.Diagnostics;

using System.Linq;

using Tensorflow.Checkpoint;

using Tensorflow.Keras.Saving.SavedModel;

using Tensorflow.Training;

using Tensorflow.Training.Saving.SavedModel;

using static Tensorflow.Binding;

namespace Tensorflow.Train

{

public abstract class Trackable: IWithTrackable

{

/// <summary>

/// Corresponding to tensorflow/python/trackable/constants.py

/// </summary>

public static class Constants

{

public static readonly string OBJECT_GRAPH_PROTO_KEY = "_CHECKPOINTABLE_OBJECT_GRAPH";

public static readonly string VARIABLE_VALUE_KEY = "VARIABLE_VALUE";

public static readonly string OBJECT_CONFIG_JSON_KEY = "OBJECT_CONFIG_JSON";

}

protected int _self_update_uid;

protected IDictionary<string, Trackable> _unconditional_dependency_names;

protected IList<TrackableReference> _unconditional_checkpoint_dependencies;

protected Dictionary<string, IList<CheckpointPosition>> _unconditional_deferred_dependencies;

protected IDictionary<string, Func<string, OneOf<BaseResourceVariable, MySaveableObject>>> _self_saveable_object_factories =

new Dictionary<string, Func<string, OneOf<BaseResourceVariable, MySaveableObject>>>();

private bool _manual_tracking = true;

private static Trackable _none = new AutoTrackable();

/// <summary>

/// This is a trick for that CSharp does not allow the key of `Dictionary` to be null.

/// The `None` can be any object that inherits `Trackable`.

/// This Property is supposed to be used only internal.

/// </summary>

public static Trackable None

{

get

{

return _none;

}

public Trackable GetTrackable()

{

return this;

}

public virtual string ObjectIdentifier

{

get => "_generic_user_object";

}

public int UpdateUid { get => _self_update_uid; set => _self_update_uid = value; }

public IList<TrackableReference> UnconditionalCheckpointDependencies { get => _unconditional_checkpoint_dependencies; }

public IDictionary<string, Trackable> UnconditionalDependencyNames { get => _unconditional_dependency_names; }

public IList<TrackableReference> CheckpointDependencies { get => UnconditionalCheckpointDependencies; }

public Dictionary<string, IList<CheckpointPosition>> DeferredDependencies => _unconditional_deferred_dependencies;

public IDictionary<string, Func<string, OneOf<BaseResourceVariable, MySaveableObject>>> SelfSaveableObjectFactories

{

get

{

return _self_saveable_object_factories;

}

set

{

_self_saveable_object_factories = value;

}

public Dictionary<string, object> CustomizedFields { get; set; } = new Dictionary<string, object>();

public virtual void SetAttr(string name, object value)

{

var t = this.GetType();

var field_info = t.GetField(name);

if(field_info is not null)

{

field_info.SetValue(this, value);

}

else

{

CustomizedFields[name] = value;

}

// On account of performance, we don't use reflection to set the attribute if it exists in `Trackable`.

// When adding new members or properties to this class, please add corresponding process to this method.

//switch (name)

//{

// case "_manual_tracking":

// {

// _manual_tracking = (bool)value;

// break;

// }

// case "_self_saveable_object_factories":

// {

// _self_saveable_object_factories = (IDictionary<string, Func<string, OneOf<BaseResourceVariable, MySaveableObject>>>)value;

// break;

// }

// case "_self_update_uid":

// {

// _self_update_uid = (int)value;

// break;

// }

// case "_unconditional_checkpoint_dependencies":

// {

// _unconditional_checkpoint_dependencies = (IList<TrackableReference>)value;

// break;

// }

// case "_unconditional_deferred_dependencies":

// {

// _unconditional_deferred_dependencies = (Dictionary<string, IList<CheckpointPosition>>)value;

// break;

// }

// case "_unconditional_dependency_names":

// {

// _unconditional_dependency_names = (IDictionary<string, Trackable>)value;

// break;

// }

// case "SelfSaveableObjectFactories":

// {

// SelfSaveableObjectFactories = (IDictionary<string, Func<string, OneOf<BaseResourceVariable, MySaveableObject>>>)value;

// break;

// }

// case "UpdateUid":

// {

// UpdateUid = (int)value;

// break;

// }

// default:

// {

// CustomizedAttributes[name] = value;

// break;

// }

}

/// <summary>

/// Restore-on-create for a variable be saved with this `Checkpointable`.

/// </summary>

/// <returns></returns>

protected virtual IVariableV1 _add_variable_with_custom_getter(VariableArgs args)

{

tf_with(ops.init_scope(), delegate

{

#pragma warning disable CS0219 // Variable is assigned but its value is never used

IInitializer checkpoint_initializer = null;

#pragma warning restore CS0219 // Variable is assigned but its value is never used

if (tf.Context.executing_eagerly())

#pragma warning disable CS0642 // Possible mistaken empty statement

;

#pragma warning restore CS0642 // Possible mistaken empty statement

else

checkpoint_initializer = null;

});

var new_variable = args.Getter(args);

// If we set an initializer and the variable processed it, tracking will not

// assign again. It will add this variable to our dependencies, and if there

// is a non-trivial restoration queued, it will handle that. This also

// handles slot variables.

if (!args.Overwrite || new_variable is RefVariable || new_variable is Trackable)

{

var res = _track_trackable(new_variable as Trackable, args.Name, args.Overwrite);

Debug.Assert(res is IVariableV1);

return res as IVariableV1;

}

else

return new_variable;

}

/// <summary>

/// Pop and load any deferred checkpoint restores into `trackable`.

/// </summary>

/// <param name="name"></param>

/// <param name="trackable"></param>

protected void _handle_deferred_dependencies(string name, IVariableV1 trackable)

{

_maybe_initialize_trackable();

// TODO

}

protected IVariableV1 _track_checkpointable(IVariableV1 checkpointable, string name, bool overwrite = false)

{

return checkpointable;

}

/// <summary>

/// Initialize dependency management.

/// </summary>

public void _maybe_initialize_trackable()

{

if(_unconditional_checkpoint_dependencies is not null)

{

return;

}

_self_update_uid = -1;

_unconditional_checkpoint_dependencies = new List<TrackableReference>();

_unconditional_dependency_names = new Dictionary<string, Trackable>();

_unconditional_deferred_dependencies = new Dictionary<string, IList<CheckpointPosition>>();

}

public virtual IDictionary<string, Trackable> _trackable_children(SaveType save_type = SaveType.CHECKPOINT,

IDictionary<string, IDictionary<Trackable, ISerializedAttributes>>? cache = null)

{

_maybe_initialize_trackable();

return _unconditional_checkpoint_dependencies.ToDictionary(x => x.Name, x => x.Refer);

}

public virtual Trackable _track_trackable(Trackable trackable, string name, bool overwrite = false)

{

_maybe_initialize_trackable();

if (!_manual_tracking) return trackable;

var new_reference = new TrackableReference(name, trackable);

var current_object = _lookup_dependency(name);

if(current_object is null)

{

_unconditional_checkpoint_dependencies.Add(new_reference);

_handle_deferred_dependencies(name, trackable);

}

_unconditional_dependency_names[name] = trackable;

return trackable;

}

/// <summary>

/// Pop and load any deferred checkpoint restores into `trackable`.

/// This method does not add a new dependency on `trackable`, but it does check if any outstanding/deferred dependencies have been queued waiting for

/// this dependency to be added (matched based on `name`). If so, `trackable` and its dependencies are restored. The restorations are

/// considered fulfilled and so are deleted.

/// `_track_trackable` is more appropriate for adding a normal/unconditional dependency, and includes handling for deferred restorations.

/// This method allows objects such as `Optimizer` to use the same restoration logic while managing conditional dependencies themselves,

/// by overriding `_checkpoint_dependencies` and `_lookup_dependency` to change the object's dependencies based on the context

/// it is saved/restored in (a single optimizer instance can have state associated with multiple graphs).

/// </summary>

/// <param name="name"></param>

/// <param name="trackable"></param>

public virtual void _handle_deferred_dependencies(string name, Trackable trackable)

{

_maybe_initialize_trackable();

trackable._maybe_initialize_trackable();

if(_unconditional_deferred_dependencies.TryGetValue(name, out var dependencies))

{

_unconditional_deferred_dependencies.Remove(name);

foreach(var checkpoint_position in dependencies.OrderByDescending(x => x.Checkpoint.RestoreUid))

{

checkpoint_position.restore(trackable);

}

// TODO(Rinne): deal with `_self_name_based_restores`

}

public virtual Trackable? _lookup_dependency(string name)

{

if (_unconditional_dependency_names.TryGetValue(name, out var dependency)) return dependency;

else return null;

}

public static Trackable convert_to_trackable(object obj, object? parent = null)

{

if (obj is Trackable)

{

return (Trackable)obj;

}

else

{

throw new NotImplementedException();

}

public virtual IDictionary<string, Trackable> deserialization_dependencies(IDictionary<string, Trackable> children)

{

return new Dictionary<string, Trackable>();

}

public virtual (IDictionary<Trackable, Trackable>, IDictionary<Tensor, Tensor>) map_resources(

SaveOptions? save_options)

{

return (new Dictionary<Trackable, Trackable>(), new Dictionary<Tensor, Tensor>());

}

public virtual List<Tensor> export_to_saved_model_graph(IDictionary<Trackable, Trackable> object_map,

IDictionary<Tensor, Tensor> tensor_map, SaveOptions? options = null)

{

var (self_object_map, self_tensor_map) = map_resources(options);

foreach (var pair in self_object_map)

{

object_map.Add(pair);

}

foreach (var pair in self_tensor_map)

{

tensor_map.Add(pair);

}

return self_tensor_map.Keys.ToList();

}

public virtual IDictionary<string, Func<string, OneOf<BaseResourceVariable, MySaveableObject>>> gather_saveables_for_checkpoint()

{

OneOf<BaseResourceVariable, MySaveableObject> create_saveable(string name = "")

{

throw new NotImplementedException();

//return new TrackableSaveable(this, null, name, null, null);

}

if (saveable_object_util.trackable_has_serialize_to_tensor(this))

{

// TODO: complete the implementation (need to complete the class `saveable_object_util.TrackableSaveable`).

Dictionary<string, Func<string, OneOf<BaseResourceVariable, MySaveableObject>>> res = new();

res[""] = create_saveable;

return res;

}

else

{

return _self_saveable_object_factories;

}

/// <summary>

/// Gathers tensors to save to the checkpoint. You should only override `serialize_to_tensors` and `restore_from_tensors`

/// if you are defining a custom resource or variable with custom ops.

/// Otherwise, please store the state of your trackable in `tf.Variable` objects

/// and add them to Trackable object hierarchy using `setattr` (for subclasses

/// of `AutoTrackable`) or overriding the `_trackable_children` method.

/// </summary>

/// <returns></returns>

/// <exception cref="NotImplementedException"></exception>

public virtual IDictionary<string, IDictionary<string, OneOf<Tensor, SaveSpec>>> serialize_to_tensors()

{

throw new NotImplementedException();

}

public virtual IDictionary<string, Operation> _restore_from_tensors(IDictionary<string, OneOf<Tensor, IDictionary<string, Tensor>>> restored_tensors)

{

throw new NotImplementedException();

}

public record class TrackableReference(string Name, Trackable Refer);

public record class SlotVariableRestoration(int OptimizerId, int SlotVariableId, string SlotName);

}

TensorFlow.NET/src/TensorFlowNET.Core/Training/Trackable.cs at master · hartoto2/TensorFlow.NET · GitHub

Expand file tree

Breadcrumbs

Trackable.cs

Latest commit

History

Breadcrumbs

Trackable.cs

File metadata and controls