Fixes in lstm

label_studio_ml/examples/yolo/tests/test_neural_nets.py

@@ -32,7 +32,7 @@ def test_multi_label_lstm():
     # Model configuration
     input_size = 50  # Number of features per time step
     output_size = 2  # Number of output classes (multi-label classification)
-    seq_len = 20  # Sequence length (number of time steps)
+    seq_len = 72  # Sequence length (number of time steps)
     hidden_size = 8  # LSTM hidden state size
 
     # Initialize device (CPU or GPU)
@@ -42,14 +42,14 @@ def test_multi_label_lstm():
     model = MultiLabelLSTM(input_size=input_size, output_size=output_size, hidden_size=hidden_size, device=device)
 
     # Example sequential data: 100 samples, each with 5 time steps and 10 features per time step
-    new_data = torch.randn(seq_len, input_size)  # Shape: (batch_size, seq_len, input_size)
-    new_labels = torch.randint(0, 2, (seq_len, output_size)).float()  # Shape: (batch_size, seq_len, output_size)
+    data = torch.randn(seq_len, input_size)  # Shape: (batch_size, seq_len, input_size)
+    labels = torch.randint(0, 2, (seq_len, output_size)).tolist() # Shape: (batch_size, seq_len, output_size)
 
     # Perform partial training with batch size of 16
-    model.partial_fit(new_data, new_labels, batch_size=16, epochs=500)
+    model.partial_fit(data, labels, batch_size=16, epochs=500)
 
     # Example prediction
-    predictions = model.predict(new_data)
+    predictions = model.predict(data)
     print(predictions)
 
     # Save the model
@@ -59,7 +59,7 @@ def test_multi_label_lstm():
     loaded_model = MultiLabelLSTM.load("lstm_model.pth")
 
     # Predict with the loaded model
-    loaded_predictions = loaded_model.predict(new_data)
-    labels = (loaded_predictions > 0.5).int()
+    loaded_predictions = loaded_model.predict(data)
+    loaded_labels = (loaded_predictions > 0.5).int()
 
-    assert torch.equal(labels, new_labels.int()), "Predicted labels do not match the training labels."
+    assert torch.equal(torch.tensor(labels), loaded_labels.int()), "Predicted labels do not match the training labels."

label_studio_ml/examples/yolo/utils/neural_nets.py

@@ -5,10 +5,12 @@
 
 from torch.utils.data import DataLoader, TensorDataset
 from torch.nn.utils.rnn import pad_sequence
+from typing import List
 
 
 logger = logging.getLogger(__name__)
 
+
 class BaseNN(nn.Module):
     def set_label_map(self, label_map):
         self.label_map = label_map
@@ -45,28 +47,49 @@ def __init__(self, input_size, output_size, device=None):
         self.device = device if device else torch.device('cpu')
         self.to(self.device)
 
-    def forward(self, x):
-        x = torch.relu(self.fc1(x))
-        x = torch.relu(self.fc2(x))
-        x = torch.sigmoid(self.output(x))  # Sigmoid for multi-label classification
-        return x
-
-    def partial_fit(self, new_data, new_labels, epochs=1):
+    def partial_fit(self, new_data, new_labels, batch_size=32, epochs=1):
         self.train()  # Set the model to training mode
+        sequence_size = self.sequence_size
 
-        # Ensure the new_data and new_labels are on the same device as the model
         new_data = torch.stack(new_data) if isinstance(new_data, list) else new_data
-        new_data = new_data.to(self.device)  # Move data to the model's device
-        new_labels = torch.tensor(new_labels, dtype=torch.float32).to(self.device)  # Move labels to the same device
+        new_labels = torch.tensor(new_labels, dtype=torch.float32)
+
+        # Split the data into small sequences by sequence_size with 1/2 overlap
+        new_data = [new_data[i:i + sequence_size] for i in range(0, len(new_data), sequence_size//2)]
+        new_data = pad_sequence(new_data, batch_first=True, padding_value=0)
+        new_labels = [new_labels[i:i + sequence_size] for i in range(0, len(new_labels), sequence_size//2)]
+        new_labels = pad_sequence(new_labels, batch_first=True, padding_value=0)
+
+        # Create a DataLoader for batching the input data
         outputs = None
+        dataset = TensorDataset(new_data, torch.tensor(new_labels, dtype=torch.float32))
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
 
         for epoch in range(epochs):
-            self.optimizer.zero_grad()  # Zero out gradients from previous steps
-            outputs = self(new_data)
-            loss = self.criterion(outputs, new_labels)  # Calculate the loss
-            loss.backward()  # Backpropagation
-            self.optimizer.step()  # Update model parameters
-            print(f'Epoch {epoch + 1}, Loss: {loss.item()}')
+            epoch_loss = 0
+            for batch_data, batch_labels in dataloader:
+                # Move batch data and labels to the appropriate device
+                batch_data = batch_data.to(self.device)
+                batch_labels = batch_labels.to(self.device)
+
+                # Zero out gradients
+                self.optimizer.zero_grad()
+
+                # Forward pass
+                outputs = self(batch_data)
+
+                # Calculate loss
+                loss = self.criterion(outputs, batch_labels)
+
+                # Backpropagation
+                loss.backward()
+
+                # Update model parameters
+                self.optimizer.step()
+
+                epoch_loss += loss.item()
+
+            print(f'Epoch {epoch + 1}, Loss: {epoch_loss / len(dataloader)}')
 
         return outputs
 
@@ -131,22 +154,28 @@ def forward(self, x):
         # Output shape: (batch_size, seq_len, output_size)
         return out
 
-    def partial_fit(self, new_data, new_labels, batch_size=32, epochs=1):
-        self.train()  # Set the model to training mode
+    def preprocess_sequence(self, sequence: List[torch.Tensor], labels=None, overlap=2):
+        sequence = torch.stack(sequence) if isinstance(sequence, list) else sequence
         sequence_size = self.sequence_size
 
-        new_data = torch.stack(new_data) if isinstance(new_data, list) else new_data
-        new_labels = torch.tensor(new_labels, dtype=torch.float32)
+        # Split the data into small sequences by sequence_size with overlap
+        chunks = [sequence[i:i + sequence_size] for i in range(0, len(sequence), sequence_size // overlap)]
+        chunks = pad_sequence(chunks, batch_first=True, padding_value=0)
 
-        # Split the data into small sequences by sequence_size with 1/2 overlap
-        new_data = [new_data[i:i + sequence_size] for i in range(0, len(new_data), sequence_size//2)]
-        new_data = pad_sequence(new_data, batch_first=True, padding_value=0)
-        new_labels = [new_labels[i:i + sequence_size] for i in range(0, len(new_labels), sequence_size//2)]
-        new_labels = pad_sequence(new_labels, batch_first=True, padding_value=0)
+        if labels is not None:
+            labels = torch.tensor(labels, dtype=torch.float32)
+            labels = [labels[i:i + sequence_size] for i in range(0, len(labels), sequence_size // overlap)]
+            labels = pad_sequence(labels, batch_first=True, padding_value=0) if labels is not None else None
+
+        return chunks, labels
+
+    def partial_fit(self, sequence, labels, batch_size=32, epochs=1):
+        self.train()  # Set the model to training mode
+        batches, label_batches = self.preprocess_sequence(sequence, labels)
 
         # Create a DataLoader for batching the input data
         outputs = None
-        dataset = TensorDataset(new_data, torch.tensor(new_labels, dtype=torch.float32))
+        dataset = TensorDataset(batches, torch.tensor(label_batches, dtype=torch.float32))
         dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
 
         for epoch in range(epochs):
@@ -156,59 +185,30 @@ def partial_fit(self, new_data, new_labels, batch_size=32, epochs=1):
                 batch_data = batch_data.to(self.device)
                 batch_labels = batch_labels.to(self.device)
 
-                # Zero out gradients
                 self.optimizer.zero_grad()
-
-                # Forward pass
-                outputs = self(batch_data)
-
-                # Calculate loss
-                loss = self.criterion(outputs, batch_labels)
-
-                # Backpropagation
-                loss.backward()
-
-                # Update model parameters
-                self.optimizer.step()
+                outputs = self(batch_data)  # Forward pass
+                loss = self.criterion(outputs, batch_labels)  # Calculate loss
+                loss.backward()  # Back propagation
+                self.optimizer.step()  # Update model parameters
 
                 epoch_loss += loss.item()
 
             print(f'Epoch {epoch + 1}, Loss: {epoch_loss / len(dataloader)}')
 
         return outputs
 
-    def predict(self, new_data, threshold=None):
-        return self.predict_in_batches(new_data)
-
-        new_data = torch.stack(new_data) if isinstance(new_data, list) else new_data
-        new_data = new_data.to(self.device)  # Move data to the model's device
-
-        self.eval()  # Set the model to evaluation mode
-        with torch.no_grad():  # Disable gradient computation
-            new_data = torch.tensor(new_data, dtype=torch.float32)
-            outputs = self(new_data)
-            if threshold is None:
-                return outputs
-
-        return (outputs > threshold).int()  # Apply threshold to get binary labels
-
-    def predict_in_batches(self, sequence):
-        sequence_len = len(sequence)
-        predictions = []
-
-        # Loop over the sequence in chunks of self.sequence_size
-        for i in range(0, sequence_len, self.sequence_size):
-            # Get a chunk of the sequence
-            chunk = torch.stack(sequence[i:i + self.sequence_size])
-
-            # If the chunk is smaller than self.sequence_size, pad it with zeros
-            if len(chunk) < self.sequence_size:
-                padding = torch.zeros(self.sequence_size - len(chunk), chunk.size(1))
-                chunk = torch.cat((chunk, padding), dim=0)
-
-            # Get predictions for the chunk
-            preds = self(chunk.unsqueeze(0))  # Add batch dimension
-            predictions.append(preds)
-
-        return torch.cat(predictions, dim=1)[0]  # Concatenate all predictions
+    def predict(self, sequence):
+        """ Split sequence into chunks with sequence_size and predict by chunks.
+        Then concatenate all predictions into one sequence of labels
+        """
+        length = len(sequence)
+        if length == 0:
+            return torch.tensor([])
 
+        batches, _ = self.preprocess_sequence(sequence, overlap=1)
+        logits = self(batches)
+
+        # Concatenate batches to sequence back
+        shape = logits.shape
+        logits = torch.reshape(logits, [shape[0] * shape[1], shape[2]])
+        return logits[0: length]