From 032ecd2b204b2c75dcb9a6bc2abc6f8b51ce42c1 Mon Sep 17 00:00:00 2001 From: makseq-ubnt Date: Thu, 12 Sep 2024 16:47:48 +0100 Subject: [PATCH] Fixes in lstm --- .../examples/yolo/tests/test_neural_nets.py | 16 +- .../examples/yolo/utils/neural_nets.py | 146 +++++++++--------- 2 files changed, 81 insertions(+), 81 deletions(-) diff --git a/label_studio_ml/examples/yolo/tests/test_neural_nets.py b/label_studio_ml/examples/yolo/tests/test_neural_nets.py index 2441c967..31df358c 100644 --- a/label_studio_ml/examples/yolo/tests/test_neural_nets.py +++ b/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." diff --git a/label_studio_ml/examples/yolo/utils/neural_nets.py b/label_studio_ml/examples/yolo/utils/neural_nets.py index 0359a7c4..76bc4246 100644 --- a/label_studio_ml/examples/yolo/utils/neural_nets.py +++ b/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,20 +185,11 @@ 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() @@ -177,38 +197,18 @@ def partial_fit(self, new_data, new_labels, batch_size=32, epochs=1): 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]