Added train by epoch for Trainer and added support for texts

meshgpt_pytorch/data.py

@@ -5,6 +5,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
@@ -106,12 +107,15 @@ def custom_collate(data, pad_id = -1):
     output = []
 
     for datum in zip(*data):
-        padded = pad_sequence(datum, batch_first = True, padding_value = pad_id)
-        output.append(padded)
+        if isinstance(datum[0], Tensor):
+            padded = pad_sequence(datum, batch_first=True, padding_value=pad_id)
+            output.append(padded)
+        else:
+            output.append(list(datum)) 
 
     output = tuple(output)
 
     if is_dict:
         output = dict(zip(keys, output))
 
-    return output
+    return output

meshgpt_pytorch/meshgpt_pytorch.py

@@ -741,6 +741,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
@@ -119,6 +120,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,
@@ -197,6 +199,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)
@@ -367,7 +370,63 @@ def forward(self):
             self.wait()
 
         self.print('training complete')
+    def train(self, num_epochs, 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-autoencoder.ckpt.epoch_{epoch}_avg_loss_{avg_epoch_loss:.3f}.pt')
+
+
+        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()
 # mesh transformer trainer
 
 class MeshTransformerTrainer(Module):
@@ -387,7 +446,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,
@@ -451,6 +512,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)
@@ -576,4 +638,58 @@ def forward(self):
 
             self.wait()
 
-        self.print('training complete')
+        self.print('training complete')
+
+    def train(self, num_epochs, 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')
+
+        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()