Added train by epoch for Trainer and added support for texts

meshgpt_pytorch/data.py

@@ -6,6 +6,7 @@
 from torch.nn.utils.rnn import pad_sequence
 
 from einops import rearrange, reduce
+from torch import nn, Tensor
 
 from beartype import beartype
 from beartype.typing import Tuple, Union, Optional, Callable, Dict
@@ -114,6 +115,7 @@ def custom_collate(data, pad_id = -1):
             datum = pad_sequence(datum, batch_first = True, padding_value = pad_id)
         else:
             datum = list(datum)
+            output.append(datum)
 
         output.append(datum)
 
@@ -122,4 +124,4 @@ def custom_collate(data, pad_id = -1):
     if is_dict:
         output = dict(zip(keys, output))
 
-    return output
+    return output

meshgpt_pytorch/meshgpt_pytorch.py

@@ -793,6 +793,7 @@ def forward(
         vertices:       TensorType['b', 'nv', 3, float],
         faces:          TensorType['b', 'nf', 3, int],
         face_edges:     Optional[TensorType['b', 'e', 2, int]] = None,
+        texts: Optional[List[str]] = None,
         return_codes = False,
         return_loss_breakdown = False,
         return_recon_faces = False,

meshgpt_pytorch/trainer.py

@@ -23,7 +23,8 @@
 from meshgpt_pytorch.data import custom_collate
 
 from meshgpt_pytorch.version import __version__
-
+import matplotlib.pyplot as plt
+from tqdm import tqdm
 from meshgpt_pytorch.meshgpt_pytorch import (
     MeshAutoencoder,
     MeshTransformer
@@ -126,6 +127,7 @@ def __init__(
         accelerator_kwargs: dict = dict(),
         optimizer_kwargs: dict = dict(),
         checkpoint_every = 1000,
+        checkpoint_every_epoch: Optional[int] = None,
         checkpoint_folder = './checkpoints',
         data_kwargs: Tuple[str, ...] = ['vertices', 'faces', 'face_edges'],
         warmup_steps = 1000,
@@ -204,6 +206,7 @@ def __init__(
         self.num_train_steps = num_train_steps
         self.register_buffer('step', torch.tensor(0))
 
+        self.checkpoint_every_epoch = checkpoint_every_epoch
         self.checkpoint_every = checkpoint_every
         self.checkpoint_folder = Path(checkpoint_folder)
         self.checkpoint_folder.mkdir(exist_ok = True, parents = True)
@@ -388,7 +391,70 @@ def forward(self):
             self.wait()
 
         self.print('training complete')
+    def train(self, num_epochs, stop_at_loss = None, diplay_graph = False):
+        epoch_losses = []  # Initialize a list to store epoch losses
+        self.model.train() 
+        for epoch in range(num_epochs): 
+            total_loss = 0.0
+            num_batches = 0
+
+            progress_bar = tqdm(self.dataloader, desc=f'Epoch {epoch + 1}/{num_epochs}') 
+
+            for data in progress_bar: 
+
+                if isinstance(data, tuple): 
+                    forward_kwargs = dict(zip(self.data_kwargs, data))
 
+                elif isinstance(data, dict): 
+                    forward_kwargs = data 
+
+
+                with self.accelerator.autocast():
+                    loss = self.model(vertices = forward_kwargs['vertices'], faces= forward_kwargs['faces'])
+                    self.accelerator.backward(loss)
+
+                if exists(self.max_grad_norm):
+                    self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+
+                self.optimizer.step()
+                self.optimizer.zero_grad()
+
+                if not self.accelerator.optimizer_step_was_skipped:
+                    with self.warmup.dampening():
+                        self.scheduler.step()
+
+                current_loss = loss.item()
+                total_loss += current_loss
+                num_batches += 1
+                progress_bar.set_postfix(loss=current_loss)
+
+
+
+            avg_epoch_loss = total_loss / num_batches 
+            epoch_losses.append(avg_epoch_loss)
+            self.print(f'Epoch {epoch + 1} average loss: {avg_epoch_loss}')
+            self.wait()
+
+            if self.checkpoint_every_epoch is not None and epoch != 0 and epoch % self.checkpoint_every_epoch == 0:
+                self.save(self.checkpoint_folder / f'mesh-autoencoder.ckpt.epoch_{epoch}_avg_loss_{avg_epoch_loss:.3f}.pt')
+
+            if stop_at_loss is not None and avg_epoch_loss < stop_at_loss: 
+                self.print(f'Stopping training at epoch {epoch} with average loss {avg_epoch_loss}')
+                if self.checkpoint_every_epoch is not None:
+                    self.save(self.checkpoint_folder / f'mesh-autoencoder.ckpt.stop_at_loss_avg_loss_{avg_epoch_loss:.3f}.pt') 
+                break   
+
+
+        self.print('Training complete') 
+        if diplay_graph:
+            plt.figure(figsize=(10, 5))
+            plt.plot(range(1, num_epochs + 1), epoch_losses, marker='o')
+            plt.title('Training Loss Over Epochs')
+            plt.xlabel('Epoch')
+            plt.ylabel('Average Loss')
+            plt.grid(True)
+            plt.show()
+        return epoch_losses[-1]
 # mesh transformer trainer
 
 class MeshTransformerTrainer(Module):
@@ -408,7 +474,9 @@ def __init__(
         ema_kwargs: dict = dict(),
         accelerator_kwargs: dict = dict(),
         optimizer_kwargs: dict = dict(),
-        checkpoint_every = 1000,
+
+        checkpoint_every = 1000, 
+        checkpoint_every_epoch: Optional[int] = None,
         checkpoint_folder = './checkpoints',
         data_kwargs: Tuple[str, ...] = ['vertices', 'faces', 'face_edges'],
         warmup_steps = 1000,
@@ -472,6 +540,7 @@ def __init__(
         self.num_train_steps = num_train_steps
         self.register_buffer('step', torch.tensor(0))
 
+        self.checkpoint_every_epoch = checkpoint_every_epoch
         self.checkpoint_every = checkpoint_every
         self.checkpoint_folder = Path(checkpoint_folder)
         self.checkpoint_folder.mkdir(exist_ok = True, parents = True)
@@ -597,4 +666,65 @@ def forward(self):
 
             self.wait()
 
-        self.print('training complete')
+        self.print('training complete')
+
+    def train(self, num_epochs, stop_at_loss = None,  diplay_graph = False):
+        epoch_losses = []  # Initialize a list to store epoch losses
+        self.model.train()
+        for epoch in range(num_epochs): 
+            total_loss = 0.0
+            num_batches = 0
+
+            progress_bar = tqdm(self.dataloader, desc=f'Epoch {epoch + 1}/{num_epochs}') 
+
+            for data in progress_bar: 
+
+                if isinstance(data, tuple): 
+                    forward_kwargs = dict(zip(self.data_kwargs, data))
+
+                elif isinstance(data, dict): 
+                    forward_kwargs = data 
+
+
+                with self.accelerator.autocast():
+                    loss = self.model(**forward_kwargs)
+                    self.accelerator.backward(loss / self.grad_accum_every)
+
+                if exists(self.max_grad_norm):
+                    self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+
+                self.optimizer.step()
+                self.optimizer.zero_grad()
+
+                if not self.accelerator.optimizer_step_was_skipped:
+                    with self.warmup.dampening():
+                        self.scheduler.step()
+
+                current_loss = loss.item()
+                total_loss += current_loss
+                num_batches += 1
+                progress_bar.set_postfix(loss=current_loss)
+
+            avg_epoch_loss = total_loss / num_batches 
+            epoch_losses.append(avg_epoch_loss)
+            self.print(f'Epoch {epoch + 1} average loss: {avg_epoch_loss}')
+            self.wait() 
+            if self.checkpoint_every_epoch is not None and epoch != 0 and epoch % self.checkpoint_every_epoch == 0:
+                self.save(self.checkpoint_folder / f'mesh-transformer.ckpt.epoch_{epoch}_avg_loss_{avg_epoch_loss:.3f}.pt')
+
+            if stop_at_loss is not None and avg_epoch_loss < stop_at_loss: 
+                self.print(f'Stopping training at epoch {epoch} with average loss {avg_epoch_loss}')
+                if self.checkpoint_every_epoch is not None:
+                    self.save(self.checkpoint_folder / f'mesh-transformer.ckpt.stop_at_loss_avg_loss_{avg_epoch_loss:.3f}.pt') 
+                break   
+
+        self.print('Training complete') 
+        if diplay_graph:
+            plt.figure(figsize=(10, 5))
+            plt.plot(range(1, num_epochs + 1), epoch_losses, marker='o')
+            plt.title('Training Loss Over Epochs')
+            plt.xlabel('Epoch')
+            plt.ylabel('Average Loss')
+            plt.grid(True)
+            plt.show()
+        return epoch_losses[-1]