Commit b173b74a authored by Benjamin Vandersmissen's avatar Benjamin Vandersmissen
Browse files

Initial commit

parent 6e94da6a
# 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
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
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
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)
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)
{
"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"
}
{
"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"
}
h5py
numpy
Pillow
torch==1.11.0
torchvision==0.12.0
tqdm
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment