diff --git a/mACHINE-LEARNINGS/classes/MemeDataset.py b/mACHINE-LEARNINGS/classes/MemeDataset.py index 1fe3cea..8daaf8a 100644 --- a/mACHINE-LEARNINGS/classes/MemeDataset.py +++ b/mACHINE-LEARNINGS/classes/MemeDataset.py @@ -21,7 +21,6 @@ def __init__(self, data_dir, transform=None): for path in self.img_path: classes.add(os.path.basename(os.path.dirname(path))) self.labels = {cls: i for i, cls in enumerate(sorted(list(classes)))} - self.resize = Resize(size = (160, 160)) def __len__(self): return len(self.img_path) @@ -32,6 +31,6 @@ def __getitem__(self, idx): label = self.labels[cls] if self.transform: - return self.transform(img), torch.tensor(label) + return self.transform['train'](img), torch.tensor(label) - return self.resize(img), torch.tensor(label) \ No newline at end of file + return img, torch.tensor(label) \ No newline at end of file diff --git a/mACHINE-LEARNINGS/models/DankCNN.py b/mACHINE-LEARNINGS/models/DankCNN.py index 479b1bc..4efd450 100644 --- a/mACHINE-LEARNINGS/models/DankCNN.py +++ b/mACHINE-LEARNINGS/models/DankCNN.py @@ -1,97 +1,17 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torchvision import models class DankCNN(nn.Module): def __init__(self, dropout=False): super(DankCNN, self).__init__() - self.act = nn.ReLU() - - self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=4) - self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2) - - self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=4) - self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) - - self.conv3 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=4) - self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) #todo - - self.linear_layer1 = nn.Linear(36992, 512) - self.linear_layer2 = nn.Linear(512, 256) - self.linear_layer3 = nn.Linear(256, 128) - self.linear_layer4 = nn.Linear(128, 64) - self.linear_layer5 = nn.Linear(64, 32) - self.linear_layer6 = nn.Linear(32, 1) - - - # self.linear_layer3 = nn.Linear(4096, 1024) - # self.linear_layer4 = nn.Linear(1024, 128) - # self.linear_layer5 = nn.Linear(128, 32) - # self.linear_layer6 = nn.Linear(32, 1) - - - + # self.pretrained_model = models.resnet18(weights='IMAGENET1K_V1') + self.pretrained_model = models.resnet50(pretrained=True) + self.n_pretrained_features = self.pretrained_model.fc.in_features + self.pretrained_model.fc = nn.Linear(self.pretrained_model.fc.in_features, 1) def forward(self, x): - - x = self.conv1(x) - x = self.act(x) - x = self.pool1(x) - - x = self.conv2(x) - x = self.act(x) - x = self.pool2(x) - - x = self.conv3(x) - x = self.act(x) - x = torch.flatten(x, 1) - - x = self.linear_layer1(x) - x = self.act(x) - x = self.linear_layer2(x) - x = self.act(x) - x = self.linear_layer3(x) - x = self.act(x) - x = self.linear_layer4(x) - x = self.act(x) - x = self.linear_layer5(x) - x = self.act(x) - x = self.linear_layer6(x) - # x = self.act(x) - - return F.sigmoid(x) - - - -# # Creating a CNN-based image classifier. -class ImageClassifier(nn.Module): - def __init__(self): - super().__init__() - self.conv_layer_1 = nn.Sequential( - nn.Conv2d(3, 64, 3, padding=1), - nn.ReLU(), - nn.BatchNorm2d(64), - nn.MaxPool2d(2)) - self.conv_layer_2 = nn.Sequential( - nn.Conv2d(64, 512, 3, padding=1), - nn.ReLU(), - nn.BatchNorm2d(512), - nn.MaxPool2d(2)) - self.conv_layer_3 = nn.Sequential( - nn.Conv2d(512, 512, kernel_size=3, padding=1), - nn.ReLU(), - nn.BatchNorm2d(512), - nn.MaxPool2d(2)) - self.classifier = nn.Sequential( - nn.Flatten(), - nn.Linear(in_features=512*3*3, out_features=2)) - def forward(self, x: torch.Tensor): - x = self.conv_layer_1(x) - x = self.conv_layer_2(x) - x = self.conv_layer_3(x) - x = self.conv_layer_3(x) - x = self.conv_layer_3(x) - x = self.conv_layer_3(x) - x = self.classifier(x) - return x \ No newline at end of file + x = self.pretrained_model(x) + return F.sigmoid(x) \ No newline at end of file diff --git a/mACHINE-LEARNINGS/train.py b/mACHINE-LEARNINGS/train.py index 8b23c47..1501e25 100644 --- a/mACHINE-LEARNINGS/train.py +++ b/mACHINE-LEARNINGS/train.py @@ -5,6 +5,7 @@ from utils.performance_measure import precision_recall_f1 from models.DankCNN import DankCNN import argparse +from torchvision import transforms def execute(train_set_size, batch_size, lr, epochs, is_verbose, weight_decay): """ @@ -34,7 +35,17 @@ def execute(train_set_size, batch_size, lr, epochs, is_verbose, weight_decay): DATA LOADING - Load all data: train, test, validation """ - dataset = MemeDataset('data/', transform=None) + data_transforms = { + 'train': transforms.Compose([ + transforms.ToPILImage(), + transforms.Resize(320), + transforms.RandomResizedCrop(224), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + } + dataset = MemeDataset('data/', transform=data_transforms) train_size = int(TRAIN_SET_SIZE*len(dataset)) test_size = len(dataset) - train_size train_set, test_set = torch.utils.data.random_split(dataset, [train_size, test_size]) @@ -50,9 +61,7 @@ def execute(train_set_size, batch_size, lr, epochs, is_verbose, weight_decay): """ model = DankCNN() model = model.to(device) - # optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY) optimizer = torch.optim.RMSprop(model.parameters(), lr=.1, alpha=.99, eps=1e-08, weight_decay=0, momentum=0, centered=False) - # loss_function = torch.nn.BCELoss() loss_function = torch.nn.HingeEmbeddingLoss() """ @@ -78,8 +87,8 @@ def execute(train_set_size, batch_size, lr, epochs, is_verbose, weight_decay): acc += (label == predictions).sum()/len(label) if IS_VERBOSE: - print('Training: Epoch %d - Batch 🌭/%d: Loss: %.4f' % - (epoch+1, len(dataloaders["train"]), train_loss / (batch_num + 1))) + print('Training: Epoch %d - Batch %d/%d 🌭 Loss: %.4f' % + (epoch+1, batch_num+1, len(dataloaders["train"]), train_loss / (batch_num + 1))) n_batch = batch_num