import model
This commit is contained in:
parent
87a2663e53
commit
ba715dfce1
4
.gitignore
vendored
4
.gitignore
vendored
@ -11,6 +11,6 @@ ipython_config.py
|
|||||||
|
|
||||||
/train/
|
/train/
|
||||||
/test/
|
/test/
|
||||||
# Remove previous ipynb_checkpoints
|
|
||||||
# git rm -r .ipynb_checkpoints/
|
__pycache__/
|
||||||
|
|
||||||
|
27
test.py
27
test.py
@ -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),
|
||||||
|
Loading…
Reference in New Issue
Block a user