Initial commit

Dockerfile

+# syntax=docker/dockerfile:1
+FROM pytorch/pytorch:latest
+COPY requirements.txt /project/requirements.txt
+RUN pip3 install -r /project/requirements.txt
+COPY code /project/code
+CMD python3 /project/code/main.py -d $LTH_DATA --lr $LTH_LR -n $LTH_N -p $LTH_P -b $LTH_B -t $LTH_T -r $LTH_R --device $LTH_DEVICE

code/HDF5Dataset.py

+import torch
+from torch.utils.data import Dataset
+from torchvision.transforms import PILToTensor
+from PIL import Image
+import h5py
+
+
+class HDF5Dataset(Dataset):
+    """
+    Expects a path to a HDF5 file that contains two datasets:
+        - data : the main image data, an uint8 array of shape (nr images, width, height, nr channels)
+        - labels : the ground truth classification labels of shape (nr images, 1)
+    """
+
+    def __init__(self, filename, transform=None, as_tensor=True):
+        self.h5file = h5py.File(filename, "r")
+        self.transform = transform
+        self.as_tensor = as_tensor
+
+        self.mode = "RGB" if len(self.h5file['data'].shape) == 4 else "L"
+
+    def __len__(self):
+        return self.h5file['data'].shape[0]
+
+    def __getitem__(self, idx):
+        bla = self.h5file['data'][idx]
+        img = Image.fromarray(self.h5file['data'][idx], mode=self.mode)
+        label = self.h5file['labels'][idx]
+
+        if self.as_tensor:
+            img = PILToTensor()(img).to(torch.float32) / 255
+
+        if self.transform:
+            img = self.transform(img)
+        return img, label

code/lenet.py

+import copy
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from pruning import add_pruning_info, can_be_pruned
+
+class LeNet5(nn.Module):
+    def __init__(self):
+        super(LeNet5, self).__init__()
+        self.conv1 = add_pruning_info(nn.Conv2d(1, 6, 5, padding=2))
+        self.pool1 = nn.AvgPool2d(2)
+        self.conv2 = add_pruning_info(nn.Conv2d(6, 16, 5))
+        self.pool2 = nn.AvgPool2d(2)
+        self.conv3 = add_pruning_info(nn.Conv2d(16, 120, 5))
+        self.fc1 = add_pruning_info(nn.Linear(120, 84))
+        self.fc2 = nn.Linear(84, 10)
+
+    def forward(self, x):
+        x = self.conv1(x).tanh()
+        x = self.pool1(x)
+        x = self.conv2(x).tanh()
+        x = self.pool2(x)
+        x = self.conv3(x).tanh()
+        x = torch.flatten(x, start_dim=1)
+        x = self.fc1(x).tanh()
+        x = self.fc2(x)
+        return x
+
+    def prob(self, x):
+        return F.softmax(self.forward(x), dim=1)
+
+    def keep_pruned_weights(self):
+        for layer in self.modules():
+            if can_be_pruned(layer):
+                with torch.no_grad():
+                    layer.weight[layer.pruning_mask] = 0

code/main.py

+from HDF5Dataset import *
+from lenet import LeNet5
+from torch.utils.data import DataLoader
+import torch.nn as nn
+from tqdm import tqdm
+from pruning import *
+import argparse
+from datetime import datetime
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-d', '--data', default='fashion', type=str, help="The dataset to use")
+parser.add_argument('--lr', default=0.0005, type=float, help="Which learning rate to use")
+parser.add_argument('-n', '--pruningepochs', default=10, type=int, help="How many times we prune the network")
+parser.add_argument('-p', '--pruningpercentage', default=0.2, type=float, help="How much do we prune each step")
+parser.add_argument('-b', '--batch', default=32, type=int, help="Batch size")
+parser.add_argument('-t', '--trainepochs', default=10, type=int, help="How many epochs do we train")
+parser.add_argument('-r', '--random', default=42, type=int, help="The random seed used")
+parser.add_argument('--device', default='cuda:0', type=str, help="The device to run on")
+args = parser.parse_args()
+
+
+traindata = HDF5Dataset("hdf5/{}/train.hdf5".format(args.data))
+testdata = HDF5Dataset("hdf5/{}/test.hdf5".format(args.data))
+
+torch.manual_seed(args.random)
+batch_size = args.batch
+lr = args.lr
+
+device = args.device if torch.cuda.is_available() else 'cpu'
+
+model = LeNet5().to(device)
+optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+criterion = nn.CrossEntropyLoss()
+
+trainloader = DataLoader(traindata, batch_size, shuffle=True)
+testloader = DataLoader(testdata, batch_size, shuffle=False)
+
+basedir = "weights/"+str(int(datetime.now().timestamp()))
+os.makedirs(basedir, exist_ok=True)
+
+
+def iterate_over_dataset(dataloader, model, criterion, optimizer, device, train=True):
+    model.train() if train else model.eval()
+    running_loss = 0
+    running_accuracy = 0
+
+    for X, y_true in tqdm(dataloader):
+        optimizer.zero_grad()
+
+        X = X.to(device)
+        y_true = y_true.to(device)
+
+        # Forward pass
+        y_hat = model(X)
+        loss = criterion(y_hat, y_true)
+        running_loss += loss.item() * X.size(0)
+
+        running_accuracy += torch.eq(torch.argmax(y_hat, dim=1), y_true).sum().detach().cpu().numpy()
+
+        # Backward pass
+        if train:
+            loss.backward()
+            optimizer.step()
+            model.keep_pruned_weights()
+
+    epoch_loss = running_loss / len(dataloader.dataset)
+    epoch_accuracy = running_accuracy / len(dataloader.dataset)
+    return model, optimizer, epoch_loss, epoch_accuracy
+
+
+def training_loop(model, criterion, optimizer, train_loader, valid_loader, epochs, device):
+    """
+    Function defining the entire training loop
+    """
+
+    # set objects for storing metrics
+    train_losses = []
+    valid_losses = []
+    train_accuracies = []
+    valid_accuracies = []
+
+    # Train model
+    for epoch in range(0, epochs):
+
+        print("Training: epoch {} / {}".format(epoch+1, epochs))
+        # training
+        model, optimizer, train_loss, train_accuracy = iterate_over_dataset(train_loader, model, criterion, optimizer, device)
+        train_losses.append(train_loss)
+        train_accuracies.append(train_accuracy)
+
+        print("Validating: epoch {} / {}".format(epoch+1, epochs))
+        # validation
+        with torch.no_grad():
+            model, optimizer, valid_loss, valid_accuracy = iterate_over_dataset(valid_loader, model, criterion, optimizer, device, train=False)
+            valid_losses.append(valid_loss)
+            valid_accuracies.append(valid_accuracy)
+
+    return model, optimizer, (train_losses, valid_losses), (train_accuracies, valid_accuracies)
+
+
+for i in range(args.pruningepochs):
+    model, _, losses, accuracies = training_loop(model, criterion, optimizer, trainloader, testloader, args.trainepochs, device)
+    print("Losses: {}".format(losses))
+    print("Accuracies: {}".format(accuracies))
+    print("Pruning % : {}".format(pruned_percentage(model)))
+
+    save_weights_pruning_masks(model, i, basedir)
+    prune_by_magnitude(model)
+    reset_init_weights(model)
+
+save_weights_pruning_masks(model, 10, basedir)
+

code/pruning.py

+import torch
+import copy
+import os
+import numpy as np
+
+
+def add_pruning_info(module):
+    assert hasattr(module, 'weight')
+    module.pruning_mask = torch.zeros_like(module.weight, dtype=torch.bool)
+    module.init_weight = copy.deepcopy(module.weight)
+    return module
+
+
+def can_be_pruned(module):
+    return hasattr(module, 'pruning_mask')
+
+
+def prune_by_magnitude(model: torch.nn.Module, percentage=0.2):
+    # TODO: prune on magnitude (absolute), calculate quantile without pruned values
+    weights = None
+    for layer in model.modules():
+        if can_be_pruned(layer):
+            reshapen_weight = torch.abs(layer.weight[torch.logical_not(layer.pruning_mask)]).reshape(-1)
+            if weights is None:
+                weights = reshapen_weight
+            else:
+                weights = torch.cat((weights, reshapen_weight))
+
+    quantile = torch.quantile(weights, percentage)
+
+    for layer in model.modules():
+
+        if can_be_pruned(layer):
+            device = layer.weight.device
+            mask = torch.lt(torch.abs(layer.weight), quantile)  # 0 means that the weight is active, 1 that it is inactive
+            layer.pruning_mask = torch.logical_or(layer.pruning_mask.to(device), mask.to(device))
+
+
+def reset_init_weights(model: torch.nn.Module):
+    for layer in model.modules():
+        if can_be_pruned(layer):
+            layer.weight = torch.nn.Parameter(layer.init_weight * torch.logical_not(layer.pruning_mask))  # reset to init weights, EXCEPT FOR pruned weights.
+
+
+def pruned_percentage(model: torch.nn.Module):
+    prunable_weights = 0
+    pruned_weights = 0
+    for layer in model.modules():
+        if can_be_pruned(layer):
+            prunable_weights += layer.weight.numel()
+            pruned_weights += layer.pruning_mask.sum()
+
+    return pruned_weights/prunable_weights
+
+
+def pruned_overall_percentage(model: torch.nn.Module):
+    all_weights = 0
+    pruned_weights = 0
+
+    for layer in model.modules():
+        if hasattr(layer, 'weight'):
+            all_weights += layer.weight.numel()
+            if can_be_pruned(layer):
+                pruned_weights += layer.pruning_mask.sum()
+
+    return pruned_weights/all_weights
+
+
+def save_weights_pruning_masks(model: torch.nn.Module, iteration=0, prefix="weights"):
+    path = prefix + "/" + str(iteration) + "/"
+    os.makedirs(path, exist_ok=True)
+
+    for name, layer in model.named_modules():
+        if can_be_pruned(layer):
+            weight = layer.weight.detach().cpu().numpy()
+            mask = layer.pruning_mask.detach().cpu().numpy()
+            np.save("{}/{}_weight.npy".format(path, name), weight)
+            np.save("{}/{}_mask.npy".format(path, name), mask)

jobRequest.json

+{
+    "name": "LTH test",
+    "deploymentEnvironment": "production",
+    "request": {
+        "resources": {
+            "cpus": 1,
+            "gpus": 1,
+            "cpuMemoryGb": 16,
+            "minCudaVersion": 11,
+            "clusterId": 7
+        },
+        "docker": {
+            "image": "gitlab+deploy-token-535:19ESJYmR2qaQn5toyB8s@gitlab.ilabt.imec.be:4567/sparse-representation-learning/lth:latest",
+            "environment":{
+                "LTH_DATA": "fashion",
+                "LTH_LR": 0.0005,
+                "LTH_N": 20,
+                "LTH_P": 0.2,
+                "LTH_B": 128,
+                "LTH_T":15,
+                "LTH_R":42,
+                "LTH_DEVICE":"cuda:0"
+            },
+            "storage": [
+                {
+                    "containerPath": "/project",
+                    "hostPath": "/project_antwerp"
+                }
+            ]
+        }
+    },
+    "description": "testing the LTH project"
+}

jobRequest_rsync.json

+{
+    "name": "LTH test",
+    "deploymentEnvironment": "production",
+    "request": {
+        "resources": {
+            "cpus": 1,
+            "gpus": 0,
+            "cpuMemoryGb": 4,
+            "minCudaVersion": 11,
+            "clusterId": 7
+        },
+        "docker": {
+            "image": "gitlab+deploy-token-535:19ESJYmR2qaQn5toyB8s@gitlab.ilabt.imec.be:4567/sparse-representation-learning/lth:latest",
+            "environment":{
+                "LTH_DATA": "fashion",
+                "LTH_LR": 0.0005,
+                "LTH_N": 10,
+                "LTH_P": 0.2,
+                "LTH_B": 32,
+                "LTH_T":10,
+                "LTH_R":42,
+                "LTH_DEVICE":"cuda:0"
+            },
+            "storage": [
+                {
+                    "containerPath": "/project",
+                    "hostPath": "/project_antwerp"
+                }
+            ],
+            "command": "bash -c 'apt-get update; apt-get install -y rsync vim; sleep 3600'"
+        }
+    },
+    "description": "transfering files to the LTH project"
+}

requirements.txt

+h5py
+numpy
+Pillow
+torch==1.11.0
+torchvision==0.12.0
+tqdm
\ No newline at end of file