Model development (#20)

ioannispol · web-flow · commit 1740a6f677eb · 2023-11-07T12:46:35.000Z
* add model blocks and tests

* add model and tests

* add torch to setup

* add torch to requirements

* update workflow

* add data loader and tests

* add opencv to requirements

* add opencv to requirements and setup

* update opencv version

* implement attention mechanism in unet

* add network visualisations

* update readme

* add visualisations

* update requirements

* devcontainer tweaks

* tweak devcontainer

* minor changes to setup.sh- add safe directory

* add utils func to make images to npy

* update dockerfile

* test training

* complete training
diff --git a/.gitignore b/.gitignore
@@ -162,3 +162,8 @@ cython_debug/
 # Datasets
 data/
 *.npy
+
+# Experiments
+./experiment/
+*.pth
+./wandb/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -7,11 +7,12 @@ repos:
     -   id: check-yaml
     -   id: check-added-large-files
 
--   repo: https://github.com/psf/black
-    rev: 21.9b0  # Use the version you prefer
-    hooks:
-    -   id: black
-        args: ['--safe']
+# -   repo: https://github.com/psf/black
+#     rev: 21.9b0  # Use the version you prefer
+#     hooks:
+#     -   id: black
+#         args: ['--safe']
+
 
 -   repo: https://github.com/pycqa/flake8
     rev: 3.9.2  # Use the version you prefer
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -0,0 +1,16 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: Current File",
+            "type": "python",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal",
+            "justMyCode": true
+        }
+    ]
+}
diff --git a/Dockerfile b/Dockerfile
@@ -8,6 +8,8 @@ RUN apt clean && \
     rm -rf /var/lib/apt/lists/* && \
     apt-get update && \
     apt-get install -y wget  \
+   # apt-get install -y libgl1-mesa-glx \
+
     git -y && \
     wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh && \
      /bin/bash ~/miniconda.sh -b -p /opt/miniconda
diff --git a/notebooks/test_unet.ipynb b/notebooks/test_unet.ipynb
@@ -122,7 +122,8 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.8.18"
+
   }
  },
  "nbformat": 4,
diff --git a/scripts/dataset.zip b/scripts/dataset.zip
diff --git a/scripts/download_data.sh b/scripts/download_data.sh
@@ -0,0 +1,27 @@
+#!/bin/bash
+
+# Check if an argument was provided
+if [ "$#" -ne 1 ]; then
+    echo "Usage: $0 YOUR_ONEDRIVE_SHARE_LINK"
+    exit 1
+fi
+
+# The OneDrive shared link is the first argument
+onedrive_link="$1"
+
+# Convert the OneDrive share link to a direct download link
+direct_link=$(echo $onedrive_link | sed 's/redir?/download?/')
+
+# Set the name of the file to download to
+output_filename="dataset.zip"
+
+# Use wget or curl to download the file
+# Uncomment the line for the tool you have available
+
+# Using wget
+wget -O "$output_filename" "$direct_link"
+
+# Using curl
+# curl -L "$direct_link" -o "$output_filename"
+
+echo "Download complete: $output_filename"
diff --git a/scripts/git_update.sh b/scripts/git_update.sh
@@ -0,0 +1,43 @@
+#!/bin/bash
+
+# Navigate to your repository directory
+# cd /path/to/your/repo
+
+# Function to set git user configuration
+setup_git_config() {
+  echo "Configuring git user information..."
+  echo "Enter your git user email:"
+  read git_email
+  git config --global user.email "$git_email"
+
+  echo "Enter your git user name:"
+  read git_name
+  git config --global user.name "$git_name"
+}
+
+# Check if git user is set, if not set it
+git_user_name=$(git config --global user.name)
+git_user_email=$(git config --global user.email)
+
+if [[ -z "$git_user_name" || -z "$git_user_email" ]]; then
+  setup_git_config
+fi
+
+# Check git status
+git status
+
+# Add changes to the index
+git add .
+
+# Prompt user for a commit message
+echo "Enter commit message:"
+read commit_message
+
+# Commit changes
+git commit -m "$commit_message"
+
+# Push changes to the remote repository
+echo "Pushing to remote repository..."
+git push
+
+echo "Git operations completed."
diff --git a/setup.py b/setup.py
@@ -13,6 +13,7 @@
         'matplotlib',
         'pandas',
         'torch ~= 2.0',
+        'torchvision',
         'opencv-python <=4.8.0.74'
     ],
     extras_require={
@@ -26,5 +27,5 @@
     packages=find_packages(
         include=['underwater_unet', 'underwater_unet*'],
         exclude=['tests', 'tests.*', 'notebooks']
-        ),
+    ),
 )
diff --git a/test.jpg b/test.jpg
diff --git a/test.pt b/test.pt
diff --git a/train.py b/train.py
@@ -1,52 +1,147 @@
+import os
+import argparse
+
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.utils.data import DataLoader
+import torch.nn.functional as F
+from torch.cuda.amp import GradScaler, autocast
+from torch.utils.data import DataLoader, random_split
 from torchvision import transforms
+from tqdm import tqdm
 
 from underwater_unet.model import UNet
 from utils.data_load import UnderwaterDataset
+from utils.dice_score import dice_loss
+
+
+def train_model(
+        model,
+        device,
+        epochs=10,
+        batch_size=1,
+        learning_rate=0.001,
+        val_percent=10.0,
+        save_checkpoint=True,
+        weight_decay=1e-8,
+        momentum=0.9,
+        amp=False,  # Add AMP flag to the function arguments
+        gradient_clipping=1.0
+):
+    # Define transform
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        # Add any other transformations here
+    ])
+
+    # Load dataset and dataloader
+    train_dataset = UnderwaterDataset(images_dir='data/images', mask_dir='data/masks', augmentations=transform)
+    n_val = int(len(train_dataset) * val_percent / 100)
+    n_train = len(train_dataset) - n_val
+    train_set, val_set = random_split(train_dataset, [n_train, n_val])
+
+    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
+    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
+
+    # Initialize model, loss, optimizer, and AMP scaler
+    criterion = nn.CrossEntropyLoss() if model.n_classes > 1 else nn.BCEWithLogitsLoss()
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(momentum, 0.999))
+    scaler = GradScaler(enabled=amp)
+
+    # Initialize learning rate scheduler
+    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2)
+
+    # Training loop
+    for epoch in tqdm(range(epochs), desc="Epochs"):
+        model.train()
+        epoch_loss = 0
+        # Inside the training loop
+        for batch in tqdm(train_loader, desc="Training", leave=False):
+            images = batch['image'].to(device=device, dtype=torch.float32, memory_format=torch.channels_last)
+            masks = batch['mask'].to(device=device, dtype=torch.float32 if model.n_classes == 1 else torch.long)
+
+            with autocast(enabled=amp):
+                masks_pred = model(images)
+                if model.n_classes == 1:
+                    masks_pred = masks_pred.squeeze(1)  # Remove channel dim because we have a single class
+                    masks = masks.squeeze(1)  # Squeeze the masks for binary case
+                    loss = criterion(masks_pred, masks)
+                    loss += dice_loss(torch.sigmoid(masks_pred), masks, multiclass=False)
+                else:
+                    # No need to squeeze masks for multi-class since CrossEntropy expects a 1D target
+                    loss = criterion(masks_pred, masks)
+                    loss += dice_loss(
+                        F.softmax(masks_pred, dim=1),
+                        F.one_hot(masks, model.n_classes).permute(0, 3, 1, 2),
+                        multiclass=True
+                    )
+            # Backward and optimization steps remain the same
+
+            assert masks_pred.shape == \
+                masks.shape, f"Output shape {masks_pred.shape} and mask shape {masks.shape} do not match"
+            # Backward pass and optimization
+            optimizer.zero_grad()
+            scaler.scale(loss).backward()
+
+            # Gradient clipping
+            if gradient_clipping > 0:
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), gradient_clipping)
+
+            scaler.step(optimizer)
+            scaler.update()
+
+            epoch_loss += loss.item()
+
+        val_loss = 0
+        model.eval()
+
+        with torch.no_grad():
+            for batch in tqdm(val_loader, desc="Validation", leave=False):
+                images = batch['image'].to(device)
+                masks = batch['mask'].to(device).unsqueeze(1).float()
+                outputs = model(images)
+                probs = torch.sigmoid(outputs)  # Apply sigmoid for dice_loss
+                dice = dice_loss(probs, masks, multiclass=False)
+                bce = criterion(outputs, masks)
+                loss = bce + dice  # Consistent with training loss
+                val_loss += loss.item()
 
+        scheduler.step(val_loss)
 
-# Hyperparameters and setup
-num_epochs = 10
-learning_rate = 0.001
-batch_size = 16
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        print(f'Epoch {epoch + 1}/{epochs}, Training Loss: {epoch_loss / len(train_loader):.4f}, \
+              Validation Loss: {val_loss / len(val_loader):.4f}')
 
-# Load dataset and dataloader
-train_dataset = UnderwaterDataset(images_dir='data/images', mask_dir='data/masks', resize_to=None)
-train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+        # Save the model
+        if save_checkpoint:
+            os.makedirs("experiment", exist_ok=True)
+            torch.save(model.state_dict(), f"experiment/model_epoch_{epoch + 1}.pth")
 
-# Load dataset and dataloader
-transform = transforms.Compose([transforms.ToTensor()])
-train_dataset = UnderwaterDataset(root_dir='data', transform=transform)
-train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    print("Training completed.")
 
-# Initialize model, loss and optimizer
-model = UNet(n_channels=3, n_classes=1).to(device)
-criterion = nn.BCEWithLogitsLoss()
-optimizer = optim.Adam(model.parameters(), lr=learning_rate)
 
-# Training loop
-for epoch in range(num_epochs):
-    model.train()
-    for batch in train_loader:
-        images = batch['image'].to(device)
-        masks = batch['mask'].to(device)
+# Rest of the script remains unchanged
 
-        # Forward pass
-        outputs = model(images)
-        loss = criterion(outputs, masks)
 
-        # Backward pass and optimization
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
+def get_args():
+    parser = argparse.ArgumentParser(description='Train the UNet on images and target masks')
+    parser.add_argument('--epochs', '-e', metavar='E', type=int, default=10, help='Number of epochs')
+    parser.add_argument('--batch-size', '-b', dest='batch_size', metavar='B', type=int, default=1, help='Batch size')
+    parser.add_argument('--learning-rate', '-l', metavar='LR', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--val-percent', '-v', dest='val_percent', type=float, default=10.0,
+                        help='Percent of the data that is used as validation (0-100)')
+    parser.add_argument('--save-checkpoint', action='store_true', help='Save a checkpoint at the end of each epoch')
+    parser.add_argument('--amp', action='store_true', help='Use automatic mixed precision for training')
+    parser.add_argument('--weight-decay', type=float, default=1e-8, help='Weight decay used in optimization')
+    parser.add_argument('--momentum', type=float, default=0.9, help='Momentum used in optimization')
+    parser.add_argument('--n-classes', type=int, default=1, help='Number of classes for segmentation')
+    parser.add_argument('--bilinear', action='store_true', help='Use bilinear upsampling')
+    # You can add more arguments as required
 
-    print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")
+    return parser.parse_args()
 
-    # Save the model
-    torch.save(model.state_dict(), f"experiment/model_epoch_{epoch + 1}.pth")
 
-print("Training completed.")
+if __name__ == '__main__':
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = UNet(n_channels=3, n_classes=1).to(device)
+    train_model(model=model, device=device)
diff --git a/underwater_unet/model.py b/underwater_unet/model.py
@@ -29,11 +29,11 @@ def __init__(self, n_channels: int, n_classes, bilinear=False) -> None:
         self.outc = (OutConv(64, n_classes))
 
     def forward(self, x):
-        x1 = checkpoint(self.inc, x)
-        x2 = checkpoint(self.down1, x1)
-        x3 = checkpoint(self.down2, x2)
-        x4 = checkpoint(self.down3, x3)
-        x5 = checkpoint(self.down4, x4)
+        x1 = checkpoint(self.inc, x, use_reentrant=False)
+        x2 = checkpoint(self.down1, x1, use_reentrant=False)
+        x3 = checkpoint(self.down2, x2, use_reentrant=False)
+        x4 = checkpoint(self.down3, x3, use_reentrant=False)
+        x5 = checkpoint(self.down4, x4, use_reentrant=False)
         x = self.up1(x5, x4)
         x = self.up2(x, x3)
         x = self.up3(x, x2)
@@ -54,7 +54,7 @@ def __init__(self, n_channels, n_classes, bilinear=False):
 
     def forward(self, x):
         # Encoder
-        x1 = checkpoint(self.inc, x)
+        x1 = checkpoint(self.inc, x,  use_reentrant=False)
         x2 = checkpoint(self.down1, x1)
         x3 = checkpoint(self.down2, x2)
         x4 = checkpoint(self.down3, x3)
diff --git a/utils/data_load.py b/utils/data_load.py
@@ -23,8 +23,10 @@ def load_image(filename: Union[str, Path]) -> Union[np.ndarray, None]:
 
 
 def mask_values_unique(idx: int, mask_dir: Path, mask_suffix: str) -> np.ndarray:
-    mask_file = list(mask_dir.glob(idx + mask_suffix + '.*'))[0]
-    mask = np.asarray(load_image(mask_file))
+    mask_file = list(mask_dir.glob('mask-' + idx + mask_suffix + '.*'))
+    if not mask_file:
+        raise FileNotFoundError(f"No mask file found for index {idx} with suffix '{mask_suffix}' in {mask_dir}")
+    mask = np.asarray(load_image(mask_file[0]))
     if mask.ndim == 2:
         return np.unique(mask)
     elif mask.ndim == 3:
@@ -99,7 +101,7 @@ def preprocess(mask_values, img: np.ndarray, resize_to: Tuple[int, int], is_mask
 
     def __getitem__(self, idx: int) -> dict:
         name = self.ids[idx]
-        mask_file = list(self.mask_dir.glob(name + self.mask_suffix + '.*'))
+        mask_file = list(self.mask_dir.glob('mask-' + name + self.mask_suffix + '.*'))
         img_file = list(self.images_dir.glob(name + '.*'))
 
         assert len(img_file) == 1, f'Either no image or multiple images found for the ID {name}: {img_file}'
