import model

This commit is contained in:
filnow 2024-04-28 19:10:12 +02:00
parent 87a2663e53
commit ba715dfce1
2 changed files with 3 additions and 28 deletions

4
.gitignore vendored
View File

@ -11,6 +11,6 @@ ipython_config.py
/train/ /train/
/test/ /test/
# Remove previous ipynb_checkpoints
# git rm -r .ipynb_checkpoints/ __pycache__/

27
test.py
View File

@ -1,36 +1,11 @@
import csv import csv
import torch import torch
import torch.nn.functional as F
from torch import nn
from torchvision import transforms, datasets from torchvision import transforms, datasets
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
from train import Model
class Model(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 32, 3, 1)
self.batchnorm1 = nn.BatchNorm2d(32)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.batchnorm2 = nn.BatchNorm2d(64)
self.conv3 = nn.Conv2d(64, 128, 3, 1)
self.fc1 = nn.Linear(128*26*26, 128)
self.fc2 = nn.Linear(128, 2)
def forward(self, x):
x = F.relu(self.batchnorm1(self.conv1(x)))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.batchnorm2(self.conv2(x)))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.conv3(x))
x = F.max_pool2d(x, 2, 2)
x = x.view(-1, 128*26*26)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.log_softmax(x, dim=1)
def get_data(IMG_SIZE: int, BATCH_SIZE: int): def get_data(IMG_SIZE: int, BATCH_SIZE: int):
testTransformer = transforms.Compose([ testTransformer = transforms.Compose([
transforms.Resize(size = (IMG_SIZE, IMG_SIZE), antialias = True), transforms.Resize(size = (IMG_SIZE, IMG_SIZE), antialias = True),