speechbrain/speechbrain/core.py at develop · speechbrain/speechbrain · GitHub

History

1491 lines (1295 loc) · 57 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

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

548

549

550

551

552

553

554

555

556

557

558

559

560

561

562

563

564

565

566

567

568

569

570

571

572

573

574

575

576

577

578

579

580

581

582

583

584

585

586

587

588

589

590

591

592

593

594

595

596

597

598

599

600

601

602

603

604

605

606

607

608

609

610

611

612

613

614

615

616

617

618

619

620

621

622

623

624

625

626

627

628

629

630

631

632

633

634

635

636

637

638

639

640

641

642

643

644

645

646

647

648

649

650

651

652

653

654

655

656

657

658

659

660

661

662

663

664

665

666

667

668

669

670

671

672

673

674

675

676

677

678

679

680

681

682

683

684

685

686

687

688

689

690

691

692

693

694

695

696

697

698

699

700

701

702

703

704

705

706

707

708

709

710

711

712

713

714

715

716

717

718

719

720

721

722

723

724

725

726

727

728

729

730

731

732

733

734

735

736

737

738

739

740

741

742

743

744

745

746

747

748

749

750

751

752

753

754

755

756

757

758

759

760

761

762

763

764

765

766

767

768

769

770

771

772

773

774

775

776

777

778

779

780

781

782

783

784

785

786

787

788

789

790

791

792

793

794

795

796

797

798

799

800

801

802

803

804

805

806

807

808

809

810

811

812

813

814

815

816

817

818

819

820

821

822

823

824

825

826

827

828

829

830

831

832

833

834

835

836

837

838

839

840

841

842

843

844

845

846

847

848

849

850

851

852

853

854

855

856

857

858

859

860

861

862

863

864

865

866

867

868

869

870

871

872

873

874

875

876

877

878

879

880

881

882

883

884

885

886

887

888

889

890

891

892

893

894

895

896

897

898

899

900

901

902

903

904

905

906

907

908

909

910

911

912

913

914

915

916

917

918

919

920

921

922

923

924

925

926

927

928

929

930

931

932

933

934

935

936

937

938

939

940

941

942

943

944

945

946

947

948

949

950

951

952

953

954

955

956

957

958

959

960

961

962

963

964

965

966

967

968

969

970

971

972

973

974

975

976

977

978

979

980

981

982

983

984

985

986

987

988

989

990

991

992

993

994

995

996

997

998

999

1000

"""Core SpeechBrain code for running experiments.

Authors

* Peter Plantinga 2020, 2023

* Abdel Heba 2020

* Mirco Ravanelli 2020

* Aku Rouhe 2021

* Andreas Nautsch 2022

* Sylvain de Langen 2023

* Adel Moumen 2023, 2024

"""

import inspect

import logging

import os

import pathlib

import shutil

import sys

import tempfile

import time

import warnings

from contextlib import contextmanager

from datetime import date

from enum import Enum, auto

from types import SimpleNamespace

import torch

import yaml

from hyperpyyaml import resolve_references

from packaging import version

from torch.nn import (

DataParallel as DP,

SyncBatchNorm,

)

from torch.nn.parallel import DistributedDataParallel as DDP

from torch.utils.data import DataLoader, DistributedSampler, IterableDataset

from tqdm import tqdm

import speechbrain as sb

from speechbrain.dataio.dataloader import LoopedLoader, SaveableDataLoader

from speechbrain.dataio.sampler import (

DistributedSamplerWrapper,

ReproducibleRandomSampler,

)

from speechbrain.utils.autocast import AMPConfig, TorchAutocast

from speechbrain.utils.distributed import is_distributed_initialized

from speechbrain.utils.logger import get_logger

from speechbrain.utils.optimizers import rm_vector_weight_decay

from speechbrain.utils.profiling import prepare_profiler

from speechbrain.utils.run_opts import RunOptions

sb.utils.quirks.apply_quirks()

logger = get_logger(__name__)

DEFAULT_LOG_CONFIG = os.path.dirname(os.path.abspath(__file__))

DEFAULT_LOG_CONFIG = os.path.join(DEFAULT_LOG_CONFIG, "log-config.yaml")

INTRA_EPOCH_CKPT_FLAG = "brain_intra_epoch_ckpt"

PYTHON_VERSION_MAJOR = 3

PYTHON_VERSION_MINOR = 8

def create_experiment_directory(

experiment_directory,

hyperparams_to_save=None,

overrides={},

log_config=DEFAULT_LOG_CONFIG,

save_env_desc=True,

"""Create the output folder and relevant experimental files.

Arguments

---------

experiment_directory : str

The place where the experiment directory should be created.

hyperparams_to_save : str

A filename of a yaml file representing the parameters for this

experiment. If passed, references are resolved, and the result is

written to a file in the experiment directory called "hyperparams.yaml".

overrides : dict

A mapping of replacements made in the yaml file, to save in yaml.

log_config : str

A yaml filename containing configuration options for the logger.

save_env_desc : bool

If True, an environment state description is saved to the experiment

directory, in a file called env.log in the experiment directory.

"""

try:

# all writing command must be done with the main_process

if sb.utils.distributed.if_main_process():

if not os.path.isdir(experiment_directory):

os.makedirs(experiment_directory)

# Write the parameters file

if hyperparams_to_save is not None:

hyperparams_filename = os.path.join(

experiment_directory, "hyperparams.yaml"

)

with open(hyperparams_to_save, encoding="utf-8") as f:

resolved_yaml = resolve_references(f, overrides)

with open(hyperparams_filename, "w", encoding="utf-8") as w:

print("# Generated %s from:" % date.today(), file=w)

print("# %s" % os.path.abspath(hyperparams_to_save), file=w)

print("# yamllint disable", file=w)

shutil.copyfileobj(resolved_yaml, w)

# Copy executing file to output directory

module = inspect.getmodule(inspect.currentframe().f_back)

if module is not None:

callingfile = os.path.realpath(module.__file__)

shutil.copy(callingfile, experiment_directory)

# Log exceptions to output automatically

log_file = os.path.join(experiment_directory, "log.txt")

logger_overrides = {

"handlers": {"file_handler": {"filename": log_file}}

}

sb.utils.logger.setup_logging(log_config, logger_overrides)

sys.excepthook = _logging_excepthook

# Log quirks again so that it makes it to the log file.

# Quirks are applied way earlier, before logging is properly setup,

# so this gives a chance to the user to see them, lowering surprise.

sb.utils.quirks.log_applied_quirks()

# Log beginning of experiment!

logger.info("Beginning experiment!")

logger.info(f"Experiment folder: {experiment_directory}")

# Save system description:

if save_env_desc:

description_str = sb.utils.logger.get_environment_description()

with open(

os.path.join(experiment_directory, "env.log"),

"w",

encoding="utf-8",

) as fo:

fo.write(description_str)

finally:

# wait for main_process if ddp is used

sb.utils.distributed.ddp_barrier()

def _logging_excepthook(exc_type, exc_value, exc_traceback):

"""Interrupt exception raising to log the error."""

logger.error("Exception:", exc_info=(exc_type, exc_value, exc_traceback))

class Stage(Enum):

"""Simple enum to track stage of experiments."""

TRAIN = auto()

VALID = auto()

TEST = auto()

@sb.utils.checkpoints.register_checkpoint_hooks

class Brain:

"""Brain class abstracts away the details of data loops.

The primary purpose of the `Brain` class is the implementation of

the ``fit()`` method, which iterates epochs and datasets for the

purpose of "fitting" a set of modules to a set of data.

In order to use the ``fit()`` method, one should sub-class the ``Brain``

class and override any methods for which the default behavior does not

match the use case. For a simple use case (e.g., training a single model

with a single dataset) the only methods that need to be overridden are:

* ``compute_forward()``

* ``compute_objectives()``

The example below illustrates how overriding these two methods is done.

For more complicated use cases, such as multiple modules that need to

be updated, the following methods can be overridden:

* ``fit_batch()``

* ``evaluate_batch()``

Arguments

---------

modules : dict[str, torch.nn.Module]

These modules are passed to the optimizer by default if they have

trainable parameters, and will have ``train()``/``eval()`` called on them.

opt_class : Optional[Type[torch.optim]]

A torch optimizer constructor that takes only the list of

parameters (e.g. a lambda or partial function definition). By default,

this will be passed all modules in ``modules`` at the

beginning of the ``fit()`` method. This behavior can be changed

by overriding the ``configure_optimizers()`` method.

hparams : Optional[dict]

Each key:value pair should consist of a string key and a hyperparameter

that is used within the overridden methods. These will

be accessible via an ``hparams`` attribute, using "dot" notation:

e.g., self.hparams.model(x).

run_opts : Optional[Union[RunOptions, dict]]

A set of options to change the runtime environment, see ``RunOptions`` for a list.

Typically in a script this comes from ``speechbrain.parse_args``, an alias

for ``RunOptions.from_command_line_args``. If an option is not defined here

(keep in mind that `parse_args` will inject some options by default),

then the option is also searched for in hparams (by key).

checkpointer : Optional[speechbrain.utils.checkpoints.Checkpointer]

By default, this will be used to load checkpoints, and will have the

optimizer added to continue training if interrupted.

Example

-------

>>> from torch.optim import SGD

>>> class SimpleBrain(Brain):

... def compute_forward(self, batch, stage):

... return self.modules.model(batch[0] * self.hparams.scalar)

...

... def compute_objectives(self, predictions, batch, stage):

... return torch.nn.functional.l1_loss(predictions, batch[0])

>>> model = torch.nn.Linear(in_features=10, out_features=10)

>>> brain = SimpleBrain(

... modules={"model": model},

... opt_class=lambda x: SGD(x, lr=0.1),

... hparams={"scalar": 5},

... run_opts={"device": "cpu"},

... )

>>> brain.fit(range(1), ([torch.rand(10, 10), torch.rand(10, 10)],))

"""

def __init__( # noqa: C901

self,

modules=None,

opt_class=None,

hparams=None,

run_opts=None,

checkpointer=None,

self.optimizers_dict = None

self.opt_class = opt_class

self.checkpointer = checkpointer

if isinstance(run_opts, dict):

run_opts = RunOptions.from_dictionary(run_opts)

# Check which options have been overridden. Order of priority

# is lowest: default < hparams < run_opts: highest

run_opt_defaults = RunOptions()

for arg, default in run_opt_defaults.as_dict().items():

if run_opts is not None and arg in run_opts.overridden_args:

if hparams is not None and arg in hparams:

logger.info(

f"{arg} which is specified in hparams was overridden "

+ f"by command line input to: {run_opts[arg]}"

)

setattr(self, arg, run_opts[arg])

# If any arg from run_opt_defaults exist in hparams and

# not in "run_opts" which is likely from command line

elif hparams is not None and arg in hparams:

logger.info(f"Run option {arg} from hparams is used")

setattr(self, arg, hparams[arg])

else:

setattr(self, arg, default)

# Check Python version

if not (

sys.version_info.major == PYTHON_VERSION_MAJOR

and sys.version_info.minor >= PYTHON_VERSION_MINOR

logger.warning(

"Detected Python "

+ str(sys.version_info.major)

+ "."

+ str(sys.version_info.minor)

+ ". We suggest using SpeechBrain with Python >="

+ str(PYTHON_VERSION_MAJOR)

+ "."

+ str(PYTHON_VERSION_MINOR)

)

# Assume `torchrun` was used if `RANK` and `LOCAL_RANK` are set

self.distributed_launch = (

os.environ.get("RANK") is not None

and os.environ.get("LOCAL_RANK") is not None

)

if self.data_parallel_backend and self.distributed_launch:

raise ValueError(

"To use data_parallel backend, start your script with:\n\t"

"python experiment.py hyperparams.yaml "

"--data_parallel_backend=True\n"

"To use DDP backend, start your script with:\n\t"

"torchrun [args] experiment.py hyperparams.yaml"

)

if self.ckpt_interval_minutes > 0 and self.ckpt_interval_steps > 0:

sys.exit(

"The options `ckpt_interval_minutes` and `ckpt_interval_steps` "

"are mutually exclusive. "

"Please keep only one active per experiment run."

)

# If device was not specified, then make best guess

if self.device is None:

self.device = sb.utils.distributed.infer_device()

# Set device type based on device string

if self.device == "cpu":

self.device_type = "cpu"

elif "cuda" in self.device:

self.device_type = "cuda"

# Set cuda device based on device string

try:

_, device_index = self.device.split(":")

torch.cuda.set_device(int(device_index))

except ValueError:

torch.cuda.set_device(0)

else:

self.device_type = self.device.split(":")[0]

# Checking that DataParallel use the right number of GPU

if self.data_parallel_backend and torch.cuda.device_count() == 0:

raise ValueError("You must have at least 1 GPU to use DataParallel")

# Put modules on the right device, accessible with dot notation

self.modules = torch.nn.ModuleDict(modules).to(self.device)

# The next line ensures that both tensors marked as parameters and standard tensors,

# such as those used in InputNormalization, are placed on the right device.

for module in self.modules:

if hasattr(self.modules[module], "to"):

self.modules[module] = self.modules[module].to(self.device)

# Make hyperparams available with dot notation too

if hparams is not None:

self.hparams = SimpleNamespace(**hparams)

# Checkpointer should point at a temporary directory in debug mode

if (

self.debug

and not self.debug_persistently

and self.checkpointer is not None

and hasattr(self.checkpointer, "checkpoints_dir")

tempdir = tempfile.TemporaryDirectory()

logger.info(

"Since debug mode is active, switching checkpointer "

f"output to temporary directory: {tempdir.name}"

)

self.checkpointer.checkpoints_dir = pathlib.Path(tempdir.name)

# Keep reference to tempdir as long as checkpointer exists

self.checkpointer.tempdir = tempdir

# Sampler should be handled by `make_dataloader`

# or if you provide a DataLoader directly, you can set

# this.train_sampler = your_sampler

# to have your_sampler.set_epoch() called on each epoch.

self.train_sampler = None

if self.auto_mix_prec:

logger.warning(

"The option `--auto_mix_prec` is deprecated and will be removed in the future. "

"Please use `--precision=fp16` instead."

)

self.precision = "fp16"

if self.bfloat16_mix_prec:

logger.warning(

"The option `--bfloat16_mix_prec` is deprecated and will be removed in the future. "

"Please use `--precision=bf16` instead."

)

self.precision = "bf16"

if self.device_type == "cpu" and (

self.precision == "fp16" or self.eval_precision == "fp16"

raise ValueError(

"The option `--precision` or `--eval_precision` is set to fp16. "

"This option is not yet supported on CPU. "

"Please use `--precision=bf16` or `--eval_precision=bf16` instead "

"to enable mixed precision on CPU."

)

gradscaler_enabled = (

self.precision == "fp16" and self.device_type == "cuda"

)

if self.skip_nonfinite_grads and gradscaler_enabled:

logger.warning(

"The option `skip_nonfinite_grads` will be ignored "

"because GradScaler is enabled and will automatically "

"skip nonfinite gradients."

)

logger.info(f"Gradscaler enabled: `{gradscaler_enabled}`")

logger.info(f"Using training precision: `--precision={self.precision}`")

logger.info(

f"Using evaluation precision: `--eval_precision={self.eval_precision}`"

)

if version.parse(torch.__version__) < version.parse("2.4.0"):

self.scaler = torch.cuda.amp.GradScaler(enabled=gradscaler_enabled)

else:

self.scaler = torch.GradScaler(

self.device, enabled=gradscaler_enabled

)

train_dtype = AMPConfig.from_name(self.precision).dtype

self.training_ctx = TorchAutocast(

device_type=self.device_type, dtype=train_dtype

)

eval_dtype = AMPConfig.from_name(self.eval_precision).dtype

self.evaluation_ctx = TorchAutocast(

device_type=self.device_type, dtype=eval_dtype

)

if gradscaler_enabled and self.checkpointer is not None:

self.checkpointer.add_recoverable(

"scaler", self.scaler, optional_load=True

)

# List parameter count for the user

self.print_trainable_parameters()

if self.distributed_launch:

self.rank = int(os.environ["RANK"])

if not is_distributed_initialized():

if self.rank > 0:

raise ValueError(

" ================ WARNING ==============="

"Please add sb.ddp_init_group() into your exp.py"

"To use DDP backend, start your script with:\n\t"

"torchrun [args] experiment.py hyperparams.yaml"

)

else:

logger.warning(

"To use DDP, please add "

"sb.utils.distributed.ddp_init_group() into your exp.py"

)

logger.info(

"Only the main process is alive, "

"all other subprocess were killed."

)

# Prepare iterating variables

self.avg_train_loss = 0.0

self.step = 0

self.optimizer_step = 0

# Add this class to the checkpointer for intra-epoch checkpoints

if self.checkpointer is not None:

self.checkpointer.add_recoverable("brain", self)

# Force default color for tqdm progressbar

if not self.tqdm_colored_bar:

self.tqdm_barcolor = dict.fromkeys(self.tqdm_barcolor, "")

# Profiler setup

self.profiler = None

if self.profile_training:

logger.info("Pytorch profiler has been activated.")

self.tot_prof_steps = (self.profile_steps + self.profile_warmup) - 1

self.profiler = prepare_profiler(

self.profile_warmup,

self.profile_steps,

self.hparams.output_folder,

)

self.raw_modules = (

self.modules.module

if hasattr(self.modules, "module")

else self.modules

)

def print_trainable_parameters(self):

"""Prints the number of trainable parameters in the model."""

total_trainable_params = 0

total_parameters = 0

for parameter in self.modules.parameters():

total_parameters += parameter.numel()

if parameter.requires_grad:

total_trainable_params += parameter.numel()

class_name = self.__class__.__name__

if total_parameters == 0:

logger.warning("The model has no parameters!")

logger.info(

f"{class_name} Model Statistics:\n"

f"* Total Number of Trainable Parameters: {total_trainable_params}\n"

f"* Total Number of Parameters: {total_parameters}\n"

f"* Trainable Parameters represent {0:.2f}% of the total size."

)

elif total_trainable_params == 0:

logger.warning("The model has no trainable parameters!")

formatted_total_params = sb.utils.logger.format_order_of_magnitude(

total_parameters

)

logger.info(

f"{class_name} Model Statistics:\n"

f"* Total Number of Trainable Parameters: {total_trainable_params}\n"

f"* Total Number of Parameters: {formatted_total_params}\n"

f"* Trainable Parameters represent {0:.4f}% of the total size."

)

else:

percentage_trainable = (

100 * total_trainable_params / total_parameters

)

formatted_trainable_params = (

sb.utils.logger.format_order_of_magnitude(

total_trainable_params

)

formatted_total_params = sb.utils.logger.format_order_of_magnitude(

total_parameters

)

logger.info(

f"{class_name} Model Statistics:\n"

f"* Total Number of Trainable Parameters: {formatted_trainable_params}\n"

f"* Total Number of Parameters: {formatted_total_params}\n"

f"* Trainable Parameters represent {percentage_trainable:.4f}% of the total size."

)

def compute_forward(self, batch, stage):

"""Forward pass, to be overridden by sub-classes.

Arguments

---------

batch : torch.Tensor or tensors

An element from the dataloader, including inputs for processing.

stage : Stage

The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST

Returns

-------

torch.Tensor or torch.Tensors

The outputs after all processing is complete.

Directly passed to ``compute_objectives()``.

"""

raise NotImplementedError

return

def compute_objectives(self, predictions, batch, stage):

"""Compute loss, to be overridden by sub-classes.

Arguments

---------

predictions : torch.Tensor or torch.Tensors

The output tensor or tensors to evaluate.

Comes directly from ``compute_forward()``.

batch : torch.Tensor or tensors

An element from the dataloader, including targets for comparison.

stage : Stage

The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST

Returns

-------

loss : torch.Tensor

A tensor with the computed loss.

"""

raise NotImplementedError

return

def on_stage_start(self, stage, epoch=None):

"""Gets called when a stage starts.

Useful for defining class variables used during the stage.

Arguments

---------

stage : Stage

The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST

epoch : int

The current epoch count.

"""

pass

def on_stage_end(self, stage, stage_loss, epoch=None):

"""Gets called at the end of a stage.

Useful for computing stage statistics, saving checkpoints, etc.

Arguments

---------

stage : Stage

The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST

stage_loss : float

The average loss over the completed stage.

epoch : int

The current epoch count.

"""

pass

def make_dataloader(

self, dataset, stage, ckpt_prefix="dataloader-", **loader_kwargs

"""Creates DataLoaders for Datasets.

This is used by ``fit()`` and ``evaluate()`` if they just receive

Datasets.

Alternatively, this can be called from outside the Brain subclass.

In that case, the DataLoader should be passed to ``fit()`` in place

of the dataset.

The Stage.TRAIN DataLoader is handled specially. It has extra args for

shuffle and drop_last. In DDP a DistributedSampler is created (unless

the dataset is an IterableDataset).

NOTE

----

Some important DataLoader arguments are passed via **loader_kwargs,

e.g., batch_size, num_workers, pin_memory.

NOTE

----

By default, ``evaluate()`` specifies ckpt_prefix=None to stop the test

DataLoader being added to the checkpointer. If you need to add a

recoverable after saving checkpoints (e.g., at test time, after

checkpointing the training), and still be able to recover reasonably,

you should probably specify ``allow_partial_load=True``.

Arguments

---------

dataset : Dataset

A set of data to use to create data loader. If the Dataset is a

DynamicItemDataset, PaddedBatch is used as the default collate_fn,

unless specified in loader_kwargs.

stage : Stage

The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST

ckpt_prefix : str, None

Prefix to use for SaveableDataLoader Checkpoint name. The Stage

name is added to this to create the full key. Set to None to not

save the DataLoader.

**loader_kwargs : dict

Additional keyword arguments to the DataLoader.

E.g., batch_size, num_workers, pin_memory.

Returns

-------

DataLoader for the input dataset

"""

# TRAIN stage is handled specially.

if stage == sb.Stage.TRAIN:

loader_kwargs = self._train_loader_specifics(dataset, loader_kwargs)

# This commented-out code block is useful when one can ensure

# metric reporting is DDP-valid for VALID & EVAL datasets.

# elif self.distributed_launch:

# loader_kwargs = sb.dataio.dataloader.distributed_loader_specifics(

# self.distributed_launch, self.rank, dataset, loader_kwargs

# )

dataloader = sb.dataio.dataloader.make_dataloader(

dataset, **loader_kwargs

)

if (

self.checkpointer is not None

and ckpt_prefix is not None

and (

isinstance(dataloader, SaveableDataLoader)

or isinstance(dataloader, LoopedLoader)

)

ckpt_key = ckpt_prefix + stage.name

self.checkpointer.add_recoverable(ckpt_key, dataloader)

return dataloader

def _train_loader_specifics(self, dataset, loader_kwargs):

sampler = loader_kwargs.get("sampler", None)

# Shuffling should really only matter for the train stage. Shuffling

# will also lead to more padding in batches if the order was otherwise

# sorted by length.

shuffle = loader_kwargs.get("shuffle", False)

if shuffle and not self.distributed_launch:

if sampler is not None:

raise ValueError(

"Cannot specify both shuffle=True"

"and a sampler in loader_kwargs"

)

seed = os.environ.get("SB_GLOBAL_SEED", 563375142)

sampler = ReproducibleRandomSampler(dataset, seed=seed)

self.train_sampler = sampler

loader_kwargs["sampler"] = self.train_sampler

# Delete the shuffle flag, since you cannot specify both a sampler and

# shuffling:

del loader_kwargs["shuffle"]

# Possibly make a DistributedSampler or a wrapper for some other sampler

if self.distributed_launch and not isinstance(dataset, IterableDataset):

# sort or not

if hasattr(self.hparams, "sorting"):

shuffle_ddp = (

self.hparams.sorting == "random"

) # False if 'ascending' or 'descending'

else:

shuffle_ddp = True

drop_last = loader_kwargs.get("drop_last", False)

# num_replicas arg is equal to world_size

# and retrieved automatically within

# DistributedSampler obj.

if sampler is not None:

self.train_sampler = DistributedSamplerWrapper(

sampler,

rank=self.rank,

drop_last=drop_last,

shuffle=shuffle,

)

# with DistributedSamplerWrapper, one must disable shuffling for dataloader

loader_kwargs["shuffle"] = False

loader_kwargs["sampler"] = self.train_sampler

elif loader_kwargs.get("batch_sampler") is None:

# no sampler and batch-sampler

self.train_sampler = DistributedSampler(

dataset,

rank=self.rank,

shuffle=shuffle_ddp,

drop_last=drop_last,

)

# with DistributedSamplerWrapper, one must disable shuffling for dataloader

loader_kwargs["shuffle"] = False

loader_kwargs["sampler"] = self.train_sampler

else: # batch_sampler was specified

self.train_sampler = DistributedSamplerWrapper(

loader_kwargs.get("batch_sampler", None),

rank=self.rank,

shuffle=shuffle_ddp,

)

loader_kwargs["batch_sampler"] = self.train_sampler

elif self.distributed_launch and isinstance(dataset, IterableDataset):

logger.warning(

"Cannot automatically solve distributed sampling "

"for IterableDataset."

)

return loader_kwargs

def on_fit_start(self):

"""Gets called at the beginning of ``fit()``, on multiple processes

if ``distributed_count > 0`` and backend is ddp.

Default implementation compiles the jit modules, initializes

optimizers, and loads the latest checkpoint to resume training.

"""

# Run this *after* starting all processes since jit/compiled modules

# cannot be pickled.

self._compile()

# Wrap modules with parallel backend after jit

self._wrap_distributed()

# Initialize optimizers after parameters are configured

self.init_optimizers()

# Load latest checkpoint to resume training if interrupted

if self.checkpointer is not None:

self.checkpointer.recover_if_possible()

def init_optimizers(self):

"""Called during ``on_fit_start()``, initialize optimizers

after parameters are fully configured (e.g. DDP, jit).

The default implementation of this method depends on an optimizer

class being passed at initialization that takes only a list

of parameters (e.g., a lambda or a partial function definition).

This creates a single optimizer that optimizes all trainable params.

Override this class if there are multiple optimizers.

"""

all_params = self.modules.parameters()

if self.opt_class is not None:

if self.remove_vector_weight_decay:

all_params = rm_vector_weight_decay(self.modules)

self.optimizer = self.opt_class(all_params)

self.optimizers_dict = {"opt_class": self.optimizer}

if self.checkpointer is not None:

self.checkpointer.add_recoverable("optimizer", self.optimizer)

else:

logger.info(

"No `opt_class` was provided to this Brain class, "

"skipping optimizer initialization."

)

def zero_grad(self, set_to_none=False):

"""Sets the gradients of all optimized ``torch.Tensor``s to zero

if ``set_to_none=False`` (default) or to None otherwise.

Setting gradients to None should save the memory, e.g.

during ``evaluate()`` and thus larger batch might be used.

"""

if self.optimizers_dict is not None:

for opt in self.freeze_optimizers(self.optimizers_dict).values():

opt.zero_grad(set_to_none=set_to_none)

elif self.opt_class is not None:

self.optimizer.zero_grad(set_to_none=set_to_none)

def on_evaluate_start(self, max_key=None, min_key=None):

"""Gets called at the beginning of ``evaluate()``

Default implementation loads the best-performing checkpoint for

evaluation, based on stored metrics.

Arguments

---------

max_key : str

Key to use for finding best checkpoint (higher is better).

By default, passed to ``self.checkpointer.recover_if_possible()``.

min_key : str

Key to use for finding best checkpoint (lower is better).

By default, passed to ``self.checkpointer.recover_if_possible()``.

"""

# Recover best checkpoint for evaluation

if self.checkpointer is not None:

self.checkpointer.recover_if_possible(

max_key=max_key, min_key=min_key

)

def fit_batch(self, batch):

"""Fit one batch, override to do multiple updates.

The default implementation depends on a few methods being defined

with a particular behavior:

* ``compute_forward()``

* ``compute_objectives()``

* ``optimizers_step()``

Also depends on having optimizers passed at initialization.

Arguments

---------

batch : list of torch.Tensors

Batch of data to use for training. Default implementation assumes

this batch has two elements: inputs and targets.

Returns

-------

detached loss

"""

should_step = (self.step % self.grad_accumulation_factor) == 0

self.on_fit_batch_start(batch, should_step)

with self.no_sync(not should_step):

with self.training_ctx:

outputs = self.compute_forward(batch, sb.Stage.TRAIN)

loss = self.compute_objectives(outputs, batch, sb.Stage.TRAIN)

scaled_loss = self.scaler.scale(

loss / self.grad_accumulation_factor

)

self.check_loss_isfinite(scaled_loss)

scaled_loss.backward()

if should_step:

self.optimizers_step()

self.on_fit_batch_end(batch, outputs, loss, should_step)

return loss.detach().cpu()

def check_loss_isfinite(self, loss):

"""Check if the loss is finite.

If the loss is not finite, log a helpful message and increment the `nonfinite_count`.

If the `nonfinite_count` exceeds the `--nonfinite_patience` threshold, stop the training

and raise an error.

This check is particularly useful when the loss becomes NaN or inf, while the

parameters and gradients remain finite. It helps prevent getting stuck in an

infinite loop during training.

Arguments

---------

loss : tensor

The loss tensor after ``backward()`` has been called but

before the optimizers ``step()``.

"""

if not torch.isfinite(loss):

self.nonfinite_count += 1

# Check if patience is exhausted

if self.nonfinite_count > self.nonfinite_patience:

raise ValueError(

"Loss is not finite and patience is exhausted. "

"To debug, wrap `fit()` with "

"autograd's `detect_anomaly()`, e.g.\n\nwith "

"torch.autograd.detect_anomaly():\n\tbrain.fit(...)"

)

else:

logger.warning("Patience not yet exhausted.")

def check_gradients(self):

"""Checks if the gradients are finite. If not, it will emit a warning and set them to zero."""

for param in self.modules.parameters():

if param.requires_grad and param.grad is not None:

if not torch.isfinite(param.grad).all():

param.grad = None

logger.warning(

f"Gradients {param.name} contain NaN or Inf. Setting to None."

)

def freeze_optimizers(self, optimizers):

"""By default, this method returns the passed optimizers.

Override this method if you want to freeze some optimizers

during training. To do so, return a of active optimizers.

"""

return optimizers

def optimizers_step(self):

"""Performs a step of gradient descent on the optimizers. This method is called every

``grad_accumulation_factor`` steps."""

# 1. get the valid optimizers, i.e., the ones that are not frozen during this step

if self.optimizers_dict is not None:

valid_optimizers = self.freeze_optimizers(self.optimizers_dict)

elif self.opt_class is not None:

# if valid_optimizers is not defined which could happen if a user is using an old

# init_optimizers() method, then we assume that the only valid optimizer is

# self.optimizer (which is the default behavior).

valid_optimizers = {"optimizer": self.optimizer}

else:

# Note: in some cases you might want to only compute gradients statistics and

# you do not need to call the optimizers.step() method. In this case, you can

# simply return from this method and skip the rest of the code.

return

# 2. unscale the gradients of the valid optimizers

for opt in valid_optimizers.values():

self.scaler.unscale_(opt)

# 3. clip gradients

# We are clipping this way because clipping on self.modules.parameters()

# can leads to NaN/Inf gradients norm as doing the concatenation

# of all parameters in a single vector can lead to overflow/underflow.

for opt in valid_optimizers.values():

torch.nn.utils.clip_grad_norm_(

opt.param_groups[0]["params"], self.max_grad_norm

)

# Note: no need to activate this flag if you are in fp16

# since GradScaler is automatically handling the nonfinite gradients

if not self.scaler.is_enabled() and self.skip_nonfinite_grads:

self.check_gradients()

# 4. step the valid optimizers

# If the scaler is disable, it simply calls optimizer.step()

for opt in valid_optimizers.values():

self.scaler.step(opt)

self.scaler.update()

for opt in valid_optimizers.values():

opt.zero_grad(set_to_none=True)

self.optimizer_step += 1

def on_fit_batch_start(self, batch, should_step):

"""Called at the beginning of ``fit_batch()``.

This method is not called under the AMP context manager. Do not assume

automatic casting of the input batch to a lower precision (e.g. fp16).

Arguments

---------

batch : list of torch.Tensors

Batch of data to use for training. Default implementation assumes

this batch has two elements: inputs and targets.

should_step : boolean

Whether optimizer.step() was called or not.

"""

pass

def on_fit_batch_end(self, batch, outputs, loss, should_step):

"""Called after ``fit_batch()``.

Arguments

---------

batch : list of torch.Tensors

Batch of data to use for training. Default implementation assumes

this batch has two elements: inputs and targets.

outputs : list or dictionary of torch.Tensors

Returned value of compute_forward().

loss : torch.Tensor

Returned value of compute_objectives().

should_step : boolean

Whether optimizer.step() was called or not.

"""

pass

@torch.no_grad()

def evaluate_batch(self, batch, stage):

"""Evaluate one batch, override for different procedure than train.

The default implementation depends on two methods being defined

with a particular behavior:

* ``compute_forward()``

* ``compute_objectives()``

Arguments

---------

batch : list of torch.Tensors

Batch of data to use for evaluation. Default implementation assumes

this batch has two elements: inputs and targets.

stage : Stage

View remainder of file in raw view

speechbrain/speechbrain/core.py at develop · speechbrain/speechbrain · GitHub

Uh oh!

Expand file tree

Breadcrumbs

core.py

Latest commit

History

Breadcrumbs

core.py

File metadata and controls