Commit 948674b7 authored by Benjamin Vandersmissen's avatar Benjamin Vandersmissen
Browse files

Made the project more flexible by adding a nr_classes parameter to the neural...

Made the project more flexible by adding a nr_classes parameter to the neural networks and a corresponding parameter to the HDF5Dataset
parent e395d75b
......@@ -16,6 +16,7 @@ class HDF5Dataset(Dataset):
self.h5file = h5py.File(filename, "r")
self.transform = transform
self.as_tensor = as_tensor
self.classes = self.h5file['classes'][0]
self.mode = "RGB" if len(self.h5file['data'].shape) == 4 else "L"
......@@ -23,7 +24,6 @@ class HDF5Dataset(Dataset):
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]
......
......@@ -6,6 +6,10 @@ from pruning import can_be_pruned
class CustomModel(nn.Module):
def __init__(self, nr_classes=10):
super(CustomModel, self).__init__()
self.nr_classes = nr_classes
def prob(self, x):
return F.softmax(self.forward(x), dim=1)
......
......@@ -6,15 +6,15 @@ from customModel import CustomModel
class LeNet5(CustomModel):
def __init__(self):
super(LeNet5, self).__init__()
def __init__(self, **kwargs):
super(LeNet5, self).__init__(**kwargs)
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)
self.fc2 = nn.Linear(84, self.nr_classes)
def forward(self, x):
x = self.conv1(x).tanh()
......
......@@ -11,7 +11,7 @@ from torchvision import transforms
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('--lr', default=0.0002, 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")
......@@ -34,8 +34,8 @@ else:
train_trans = None
test_trans = None
traindata = HDF5Dataset("hdf5/{}/train.hdf5".format(args.data), transform=train_trans)
testdata = HDF5Dataset("hdf5/{}/test.hdf5".format(args.data), transform=test_trans)
traindata = HDF5Dataset("../../datasets-hdf5/hdf5/{}/train.hdf5".format(args.data), transform=train_trans)
testdata = HDF5Dataset("../../datasets-hdf5/hdf5/{}/test.hdf5".format(args.data), transform=test_trans)
torch.manual_seed(args.random)
batch_size = args.batch
......@@ -43,13 +43,10 @@ lr = args.lr
device = args.device if torch.cuda.is_available() else 'cpu'
model = model_dict[args.model]().to(device)
model = model_dict[args.model](nr_classes=traindata.classes).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)
......@@ -103,17 +100,19 @@ def training_loop(model, criterion, optimizer, train_loader, valid_loader, epoch
train_losses.append(train_loss)
train_accuracies.append(train_accuracy)
print("Validating: epoch {} / {}".format(epoch+1, epochs))
# 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)
# 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):
trainloader = DataLoader(traindata, batch_size, shuffle=True)
testloader = DataLoader(testdata, batch_size, shuffle=False)
model, _, losses, accuracies = training_loop(model, criterion, optimizer, trainloader, testloader, args.trainepochs, device)
print("Losses: {}".format(losses))
print("Accuracies: {}".format(accuracies))
......
......@@ -47,6 +47,7 @@ def reinit_weights(model: torch.nn.Module):
if can_be_pruned(layer):
layer.reset_parameters()
def pruned_percentage(model: torch.nn.Module):
prunable_weights = 0
pruned_weights = 0
......
......@@ -6,8 +6,8 @@ from customModel import CustomModel
class VGG11(CustomModel):
def __init__(self, nr_classes=17):
super(VGG11, self).__init__()
def __init__(self, **kwargs):
super(VGG11, self).__init__(**kwargs)
self.conv1 = add_pruning_info(nn.Conv2d(3, 64, kernel_size=3, padding=1))
self.pool1 = nn.MaxPool2d(2)
self.conv2 = add_pruning_info(nn.Conv2d(64, 128, kernel_size=3, padding=1))
......@@ -25,7 +25,7 @@ class VGG11(CustomModel):
self.drop1 = nn.Dropout(0.5)
self.fc2 = add_pruning_info(nn.Linear(4096, 4096))
self.drop2 = nn.Dropout(0.5)
self.fc3 = add_pruning_info(nn.Linear(4096, nr_classes))
self.fc3 = add_pruning_info(nn.Linear(4096, self.nr_classes))
def forward(self, x):
x = self.conv1(x).relu()
......
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