deduplicate training loop logic

external/fv3fit/fv3fit/pytorch/graph/train.py

@@ -108,11 +108,11 @@ def train_graph_model(
     )
 
     if validation_batches is not None:
-        val_state = get_state(data=validation_batches)
+        val_state = get_state(data=validation_batches).unbatch()
     else:
         val_state = None
 
-    train_state = get_state(data=train_batches)
+    train_state = get_state(data=train_batches).unbatch()
 
     train_model = build_model(
         hyperparameters.graph_network, n_state=next(iter(train_state)).shape[-1]

external/fv3fit/fv3fit/pytorch/training_loop.py

@@ -1,6 +1,6 @@
 import dataclasses
 import logging
-from typing import Callable, Optional
+from typing import Any, Callable, Optional
 import numpy as np
 import torch
 import tensorflow_datasets as tfds
@@ -52,52 +52,24 @@ def fit_loop(
             optimizer: type of optimizer for the model
             loss_config: configuration of loss function
         """
-        train_data = train_data.shuffle(buffer_size=self.shuffle_buffer_size).batch(
-            self.samples_per_batch
+
+        def evaluate_on_batch(batch_state, model):
+            batch_input = torch.as_tensor(batch_state[0]).float().to(DEVICE)
+            batch_output = torch.as_tensor(batch_state[1]).float().to(DEVICE)
+            loss: torch.Tensor = loss_config.loss(model(batch_input), batch_output)
+            return loss
+
+        return _train_loop(
+            model=train_model,
+            train_data=train_data,
+            validation_data=validation_data,
+            evaluate_on_batch=evaluate_on_batch,
+            optimizer=optimizer,
+            n_epoch=self.n_epoch,
+            shuffle_buffer_size=self.shuffle_buffer_size,
+            samples_per_batch=self.samples_per_batch,
+            validation_batch_size=self.validation_batch_size,
         )
-        train_data = tfds.as_numpy(train_data)
-        if validation_data is not None:
-            if self.validation_batch_size is None:
-                validation_batch_size = sequence_size(validation_data)
-            else:
-                validation_batch_size = self.validation_batch_size
-            validation_data = validation_data.batch(validation_batch_size)
-            validation_data = tfds.as_numpy(validation_data)
-            min_val_loss = np.inf
-            best_weights = None
-        for i in range(1, self.n_epoch + 1):  # loop over the dataset multiple times
-            logger.info("starting epoch %d", i)
-            train_model = train_model.train()
-            train_losses = []
-            for batch_state in train_data:
-                batch_input = torch.as_tensor(batch_state[0]).float().to(DEVICE)
-                batch_output = torch.as_tensor(batch_state[1]).float().to(DEVICE)
-                optimizer.zero_grad()
-                loss: torch.Tensor = loss_config.loss(
-                    train_model(batch_input), batch_output
-                )
-                loss.backward()
-                train_losses.append(loss)
-                optimizer.step()
-            train_loss = torch.mean(torch.stack(train_losses))
-            logger.info("train loss: %f", train_loss)
-            if validation_data is not None:
-                val_model = train_model.eval()
-                val_losses = []
-                for batch_state in validation_data:
-                    batch_input = torch.as_tensor(batch_state[0]).float().to(DEVICE)
-                    batch_output = torch.as_tensor(batch_state[1]).float().to(DEVICE)
-                    with torch.no_grad():
-                        val_losses.append(
-                            loss_config.loss(val_model(batch_input), batch_output)
-                        )
-                val_loss = torch.mean(torch.stack(val_losses))
-                logger.info("val_loss %f", val_loss)
-                if val_loss < min_val_loss:
-                    min_val_loss = val_loss
-                    best_weights = train_model.state_dict()
-        if validation_data is not None:
-            train_model.load_state_dict(best_weights)
 
 
 @dataclasses.dataclass
@@ -113,10 +85,11 @@ class AutoregressiveTrainingConfig:
     """
 
     n_epoch: int = 20
-    buffer_size: int = 50_000
+    shuffle_buffer_size: int = 50_000
     samples_per_batch: int = 1
     save_path: str = "weight.pt"
     multistep: int = 1
+    validation_batch_size: Optional[int] = None
 
     def fit_loop(
         self,
@@ -136,48 +109,28 @@ def fit_loop(
             optimizer: type of optimizer for the model
             loss_config: configuration of loss function
         """
-        train_data = (
-            train_data.unbatch()
-            .shuffle(buffer_size=self.buffer_size)
-            .batch(self.samples_per_batch)
-        )
-        train_data = tfds.as_numpy(train_data)
-        if validation_data is not None:
-            validation_data = validation_data.unbatch()
-            n_validation = sequence_size(validation_data)
-            validation_state = (
-                torch.as_tensor(next(iter(validation_data.batch(n_validation))).numpy())
-                .float()
-                .to(DEVICE)
+
+        def evaluate_on_batch(batch_state, model):
+            batch_state = torch.as_tensor(batch_state).float().to(DEVICE)
+            loss: torch.Tensor = evaluate_model(
+                batch_state=batch_state,
+                model=train_model,
+                multistep=self.multistep,
+                loss=loss_config.loss,
             )
-            min_val_loss = np.inf
-        else:
-            validation_state = None
-        for _ in range(1, self.n_epoch + 1):  # loop over the dataset multiple times
-            train_model = train_model.train()
-            for batch_state in train_data:
-                batch_state = torch.as_tensor(batch_state).float().to(DEVICE)
-                optimizer.zero_grad()
-                loss = evaluate_model(
-                    batch_state=batch_state,
-                    model=train_model,
-                    multistep=self.multistep,
-                    loss=loss_config.loss,
-                )
-                loss.backward()
-                optimizer.step()
-            if validation_state is not None:
-                val_model = train_model.eval()
-                with torch.no_grad():
-                    val_loss = evaluate_model(
-                        validation_state,
-                        model=val_model,
-                        multistep=self.multistep,
-                        loss=loss_config.loss,
-                    )
-                if val_loss < min_val_loss:
-                    min_val_loss = val_loss
-                    torch.save(train_model.state_dict(), self.save_path)
+            return loss
+
+        return _train_loop(
+            model=train_model,
+            train_data=train_data,
+            validation_data=validation_data,
+            evaluate_on_batch=evaluate_on_batch,
+            optimizer=optimizer,
+            n_epoch=self.n_epoch,
+            shuffle_buffer_size=self.shuffle_buffer_size,
+            samples_per_batch=self.samples_per_batch,
+            validation_batch_size=self.validation_batch_size,
+        )
 
 
 def evaluate_model(
@@ -193,3 +146,53 @@ def evaluate_model(
         target_state = batch_state[:, step + 1, :]
         total_loss += loss(state_snapshot, target_state)
     return total_loss / multistep
+
+
+def _train_loop(
+    model: torch.nn.Module,
+    train_data: tf.data.Dataset,
+    validation_data: tf.data.Dataset,
+    evaluate_on_batch: Callable[[Any, torch.nn.Module], torch.Tensor],
+    optimizer: torch.optim.Optimizer,
+    n_epoch: int,
+    shuffle_buffer_size: int,
+    samples_per_batch: int,
+    validation_batch_size: Optional[int] = None,
+):
+
+    train_data = train_data.shuffle(buffer_size=shuffle_buffer_size).batch(
+        samples_per_batch
+    )
+    train_data = tfds.as_numpy(train_data)
+    if validation_data is not None:
+        if validation_batch_size is None:
+            validation_batch_size = sequence_size(validation_data)
+        validation_data = validation_data.batch(validation_batch_size)
+        validation_data = tfds.as_numpy(validation_data)
+        min_val_loss = np.inf
+        best_weights = None
+    for i in range(1, n_epoch + 1):  # loop over the dataset multiple times
+        logger.info("starting epoch %d", i)
+        train_model = model.train()
+        train_losses = []
+        for batch_state in train_data:
+            optimizer.zero_grad()
+            loss = evaluate_on_batch(batch_state, train_model)
+            loss.backward()
+            train_losses.append(loss)
+            optimizer.step()
+        train_loss = torch.mean(torch.stack(train_losses))
+        logger.info("train loss: %f", train_loss)
+        if validation_data is not None:
+            val_model = model.eval()
+            val_losses = []
+            for batch_state in validation_data:
+                with torch.no_grad():
+                    val_losses.append(evaluate_on_batch(batch_state, val_model))
+            val_loss = torch.mean(torch.stack(val_losses))
+            logger.info("val_loss %f", val_loss)
+            if val_loss < min_val_loss:
+                min_val_loss = val_loss
+                best_weights = train_model.state_dict()
+    if validation_data is not None:
+        train_model.load_state_dict(best_weights)