diff --git a/model.py b/model.py index ef7e257..153508f 100644 --- a/model.py +++ b/model.py @@ -4,40 +4,48 @@ class DogCatClassifier(nn.Module): def __init__ (self): super().__init__() + # 3 color (RGB) image, so tensor is of shape (B x 3 x H X W) self.conv1 = nn.Sequential( - nn.Conv2d(3, 32, 3, padding = 1), - nn.ReLU(inplace = True), - nn.MaxPool2d(2), - nn.BatchNorm2d(32) + nn.Conv2d(3, 32, 3, padding = 1), # passes conv kernel over batch and increases num channels from 3 (for RBG) to 32 + nn.ReLU(inplace = True), # relu to add nonlinearity + nn.MaxPool2d(2), # reduces h and w of img by a factor of 2 + nn.BatchNorm2d(32) #normalizes over z distribution https://arxiv.org/abs/1502.03167 ) + + # tensor size is now (B x 32 x h/2 x w/2) self.conv2 = nn.Sequential( - nn.Conv2d(32, 64, 3, padding = 1), + nn.Conv2d(32, 64, 3, padding = 1), # 32 channels to 64 ch with 3x3 kernel nn.ReLU(inplace = True), - nn.MaxPool2d(2), - nn.BatchNorm2d(64) + nn.MaxPool2d(2), # reduces # reduces h and w of img by a factor of 2 + nn.BatchNorm2d(64) #normalizes over z distribution https://arxiv.org/abs/1502.03167 ) + # tensor size is now (B x 64 x h/4 x w/4) + self.conv3 = nn.Sequential( - nn.Conv2d(64, 128, 3, padding = 1), + nn.Conv2d(64, 128, 3, padding = 1), # 64 channels to 128 ch with 3x3 kernel nn.ReLU(inplace = True), - nn.MaxPool2d(2), - nn.BatchNorm2d(128) + nn.MaxPool2d(2), # reduces # reduces h and w of img by a factor of 2 + nn.BatchNorm2d(128) # normalizes over z distribution https://arxiv.org/abs/1502.03167 ) + + # tensor size is now (B x 128 x h/8 x w/8) self.fc1 = nn.Linear(128 * 4 * 4 , 512)# 2048, lowkey had to calculator it lol - self.dropout = 0.5 # tunable - self.fc2 = nn.Linear(512, 1) + self.dropout = 0.5 # tunable, removes half of the values and replaces them with 0s + self.fc2 = nn.Linear(512, 1) # 512 ch to 1 ch output def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = self.conv3(x) x = x.view(x.size(0), -1) + # reformats for use in linear layer x = self.fc1(x) - x = nn.functional.relu(x) + x = nn.functional.relu(x) # relu to add nonlinearity x = self.fc2(x) - x = nn.functional.sigmoid(x) + x = nn.functional.sigmoid(x) # 1 / 1 + e ^(-x) return x diff --git a/test.py b/test.py index 908c00d..4225af0 100644 --- a/test.py +++ b/test.py @@ -4,23 +4,21 @@ from torch.utils.data import DataLoader from dogs_cats_ds import DogCatDataset from model import DogCatClassifier from consts import TEST_DATA -import torch.optim as optim def test(model: nn.Module, test_loader: DataLoader, criterion, device): model.eval() test_loss = 0 correct = 0 total = 0 - with torch.no_grad(): + with torch.no_grad(): # do not update gradients for img, lab in test_loader: - img, lab = img.to(device), lab.to(device).float().view(-1, 1) + img, lab = img.to(device), lab.to(device).float().view(-1, 1) # similar to how we did it in train, offset both to a gpu for better perf out = model(img) - loss = criterion(out, lab) + loss = criterion(out, lab) # evaluate loss test_loss += loss.item() - pred = (out > 0.5).float() - total += lab.size(0) - correct += (pred == lab).sum().item() - + pred = (out > 0.5).float() #same as with train, if its < 0.5 return 0 else 1 + total += lab.size(0) # total is increased by batch size + correct += (pred == lab).sum().item() # correct only += 1 if the prediction matches the label print(f'test loss: {test_loss / len(test_loader):.4f}, test_acc: {100*correct/total:.2f}%') model.train() @@ -33,6 +31,6 @@ if __name__ == "__main__": dog_test_loader = DataLoader(dog_test_dataset, batch_size = 32, shuffle = False) # since its test, bad to shuffle model = DogCatClassifier() criterion = nn.BCELoss() - model.load_state_dict(torch.load('dog_cat_classifier.pth', map_location = device, weights_only = True)) + model.load_state_dict(torch.load('dog_cat_classifier.pth', map_location = device, weights_only = True)) # loads what we trained in train.py test(model, dog_test_loader, criterion, device) diff --git a/train.py b/train.py index 1f0cd29..e710f50 100644 --- a/train.py +++ b/train.py @@ -1,33 +1,33 @@ import torch import torch.nn as nn import torch.optim as optim -from torch.utils.data import DataLoader, Dataset +from torch.utils.data import DataLoader from consts import TRAIN_DATA from tqdm import tqdm from model import DogCatClassifier from dogs_cats_ds import DogCatDataset def train(model: nn.Module, train_loader: DataLoader, criterion, optimizer, device, epochs): - model.to(device) - model.train() + model.to(device) # send to gpu if there is one, otherwise toss it over to cpu + model.train() #train mode means that all gradients are active and modifiable - for epoch in tqdm(range(epochs)): - running_loss = 0.0 + for epoch in tqdm(range(epochs)): # wrapper around for loop to add a nice progress bar + running_loss = 0.0 # start the loss, amount of cats and dogs we guess correctly, and complete samples at 0 (float 0 in case of loss since it can be a float) correct = 0 total = 0 - for i, (img, lab) in enumerate(train_loader): - img, lab = img.to(device), lab.to(device).float().view(-1, 1) - optimizer.zero_grad() + for i, (img, lab) in enumerate(train_loader): # for each image, label pair in the dataset + img, lab = img.to(device), lab.to(device).float().view(-1, 1) # send the image and label to the gpu if there is one else send to cpu, .view(-1, 1) returns the same tensor data but with the shape of the last dimension + optimizer.zero_grad() # resets gradients to zero when we initialize. - out = model(img) - loss = criterion(out, lab) - loss.backward() - optimizer.step() + out = model(img) # outputs are the results of our model on the image (sigmoid) + loss = criterion(out, lab) # loss is difference between expected and real label from prediction + loss.backward() # backprop using autograd + optimizer.step() # update optimizer - running_loss += loss.item() - pred = (out > 0.5).float() - total += lab.size(0) - correct += (pred == lab).sum().item() + running_loss += loss.item() # loss in epoch updated with loss + pred = (out > 0.5).float() # prediction is 0 if less than 0.5 else 1 + total += lab.size(0) #total samples is increased by the 0th dim of the tensor(batch size) + correct += (pred == lab).sum().item() # only add 1 to the correct count if the actual label (dog) = the predicted label(dog) if (i + 1) % 50: print(f'yo its epoch {epoch + 1} out of {epochs} and we on minibatch {i + 1} / {len(train_loader)}. Loss lookin like: {running_loss/100:.4f}, acc lookin like {100 * correct / total :.2f}%') @@ -47,6 +47,6 @@ if __name__ == "__main__": print(model) train(model = model, train_loader = dog_train_loader, criterion = criterion, optimizer = optimizer, device = device, epochs = 10) - torch.save(model.state_dict(), 'dog_cat_classifier.pth') + torch.save(model.state_dict(), 'dog_cat_classifier.pth') # saves model to pth file print('done w train, model saved')