diff --git a/src/humanoid_vision/training_pipeline.py b/src/humanoid_vision/training_pipeline.py new file mode 100644 index 0000000..396d106 --- /dev/null +++ b/src/humanoid_vision/training_pipeline.py @@ -0,0 +1,97 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader, random_split + +class TrainingPipeline (): + def __init__(self,model,data,batch_size=32,lr =0.01,epochs=100, optimizer = optim.Adam, criterion = nn.BCELoss, device = None, train_percent = 0.8): + ''' + Initialize the Training Pipeline + Args: + :param model: PASS IN MODEL CLASS. + :param data: pass into data in tensor form. + :param batch_size: batch size for training. + :param lr: learning rate. + :param epochs: Number of training epochs. + :param device: The device model and data will be passed into(ie cpu or cuda). + :param optimizer: PASS IN THE OPTIMIZER CLASS. + :param criterion: PASS IN LOSS FUNCTION CLASS. + :param train_percent: used to split into training and validation + ''' + self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.model = model.to(self.device) + self.batch_size = batch_size + self.criterion = criterion() + self.optimizer = optimizer(self.model.parameters(),lr = lr) + self.epochs = epochs + + #train test split(default 8:2) + train_size = int(train_percent* len(data)) + val_size = len(data) - train_size + train_data , val_data = random_split(data,[train_size,val_size]) + + self.train_loader = DataLoader(train_data,batch_size = batch_size,shuffle = True) + self.val_loader = DataLoader(val_data, batch_size = batch_size, shuffle = False) + + + + def train_one_epoch(self): + ''' + Runs one training loop through the model using training data + Returns: + float: Average loss for the batch + ''' + self.model.train() + total_loss = 0 + #loops through each batch + for xb,yb in self.train_loader: + xb , yb = xb.to(self.device) , yb.to(self.device) + preds = self.model(xb) + loss = self.criterion(preds,yb) + total_loss += loss.item() + + self.optimizer.zero_grad() + loss.backward() + self.optimizer.step() + + return total_loss / len(self.train_loader) #returns average loss for the batch + + + def validate(self): + ''' + Evaluates the current model using testing data + Returns: + float: average loss for the batch + ''' + self.model.eval() + total_loss = 0 + with torch.inference_mode(): + #loops through each batch + for xb,yb in self.val_loader: + xb , yb = xb.to(self.device) , yb.to(self.device) + preds = self.model(xb) + loss = self.criterion(preds,yb) + total_loss += loss.item() + + + return total_loss / len(self.val_loader) #returns average loss for the batch + + def training(self): + ''' + Loop where all the training actually happens, prints out training and validation loss + ''' + for epoch in range(self.epochs): + train_loss = self.train_one_epoch() + val_loss = self.validate() + if(epoch%5 == 0): + print(f"Train loss: {train_loss} Validation loss:{val_loss}") + + def save(self,path): + ''' + Saves the model to the given path + Args: + :param path: path where model is saved + + ''' + torch.save(self.model.state_dict(),path) + diff --git a/tests/training_pipe_test.py b/tests/training_pipe_test.py new file mode 100644 index 0000000..e52a1cf --- /dev/null +++ b/tests/training_pipe_test.py @@ -0,0 +1,19 @@ +from src.humanoid_vision.training_pipeline import TrainingPipeline +#from training_pipeline import TrainingPipeline +import torch +import torch.nn as nn +from torch.utils.data import TensorDataset + +testModel = nn.Sequential( + nn.Linear(3,1), + nn.Sigmoid() +) + +X = torch.randn(100,3) +print(X) +y = torch.randint(0, 2, (100, 1)).float() + +testData = TensorDataset(X,y) + +pipeline = TrainingPipeline(model = testModel,batch_size = 10,data = testData, device = "cpu") +pipeline.training() \ No newline at end of file