Commit 4f388d5e authored by Daniil Pastukhov's avatar Daniil Pastukhov

Final commit

parent b6164ba6
import pytorch_lightning as pl
import torch
from torch import nn, optim
import torch.nn.functional as F
import torchmetrics
def conv2d_size_out(size, kernel_size=5, stride=1):
"""
Determine the output size after applying convolution operation.
"""
return (size - (kernel_size - 1) - 1) // stride + 1
class CNN(pl.LightningModule):
def __init__(self, img_size, conv_filters, conv_kernels, conv_strides):
super().__init__()
self.fc_input_size_h = img_size
self.fc_input_size_w = img_size
# Conv layers
self.conv_layers = []
for i in range(len(conv_filters) - 1):
self.conv_layers.append(nn.Conv2d(in_channels=conv_filters[i],
out_channels=conv_filters[i + 1],
kernel_size=conv_kernels[i],
stride=conv_strides[i]))
self.conv_layers.append(nn.BatchNorm2d(conv_filters[i + 1]))
self.conv_layers.append(nn.ReLU())
self.fc_input_size_h = conv2d_size_out(self.fc_input_size_h, conv_kernels[i], 2)
self.fc_input_size_w = conv2d_size_out(self.fc_input_size_w, conv_kernels[i], 2)
self.conv_layers = nn.Sequential(*self.conv_layers)
# Fully connected layers
fc_layers = [self.fc_input_size_h * self.fc_input_size_w * conv_filters[-1], 100]
self.fc_layers = []
for i in range(len(fc_layers) - 1):
self.fc_layers.append(nn.Linear(fc_layers[i], fc_layers[i + 1]))
self.fc_layers.append(nn.ReLU() if i != len(fc_layers) - 2 else nn.LogSoftmax(dim=0))
self.fc_layers = nn.Sequential(nn.Flatten(), *self.fc_layers)
def forward(self, x):
x = self.conv_layers(x)
x = self.fc_layers(x)
return x
def configure_optimizers(self):
return optim.Adam(self.parameters(), lr=1e-3)
def common_step(self, batch, batch_id):
images, targets = batch
images = images.to(self.device).permute(0, 3, 1, 2).float()
target = targets.to(self.device).long()
output = self(images)
loss = F.nll_loss(output, target)
f1 = torchmetrics.functional.f1(output, target, average='macro', num_classes=100)
return loss, f1
def training_step(self, batch, batch_id):
loss, f1 = self.common_step(batch, batch_id)
self.log('train/loss', loss)
self.log('train/f1', f1)
return loss
@torch.no_grad()
def validation_step(self, batch, batch_id):
loss, f1 = self.common_step(batch, batch_id)
self.log('val/loss', loss)
self.log('val/f1', f1)
return loss
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
File added
from typing import Tuple
from pathlib import Path
import logging
import pandas as pd
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import WandbLogger
from skimage.io import imread
import albumentations as A
import torch
from torch.utils.data import Dataset, DataLoader
import pytorch_lightning as pl
from cnn import CNN
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
class PetDataset(Dataset):
def __init__(self, image_paths, df, transforms) -> None:
self.path_names = image_paths
self.targets = None
if 'Pawpularity' in df.columns:
self.targets = df[df.index.isin([p.stem for p in self.path_names])]['Pawpularity'].values
self.transforms = transforms
def __len__(self) -> int:
return len(self.path_names)
def __getitem__(self, index: int) -> Tuple:
image = imread(self.path_names[index])
image = self.transforms(image=image)['image']
if self.targets is None:
return image
return image, self.targets[index] - 1
def train():
LOAD_PATH = Path('data')
logger = WandbLogger('petfinder_cnn', project='petfinder', log_model='all')
train_df = pd.read_csv(Path(LOAD_PATH, 'train.csv'), index_col='Id')
test_df = pd.read_csv(Path(LOAD_PATH, 'test.csv'), index_col='Id')
train_image_paths = list(Path(LOAD_PATH, 'train').glob('*'))
test_image_paths = list(Path(LOAD_PATH, 'test').glob('*'))
img_size = 64
conv_filters = [3, 8, 16, 32, 64]
conv_kernels = [5, 5, 5, 5]
conv_strides = [2, 2, 2, 2]
train_transforms = A.Compose([
A.Resize(img_size, img_size)
])
val_transforms = A.Compose([
A.Resize(img_size, img_size)
])
train_dataset = PetDataset(train_image_paths, train_df, train_transforms)
test_dataset = PetDataset(test_image_paths, test_df, val_transforms)
val_size = int(len(train_dataset) * 0.8)
train_set, val_set = torch.utils.data.random_split(train_dataset, lengths=[len(train_dataset) - val_size, val_size])
batch_size = 128
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=0)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
model = CNN(img_size, conv_filters, conv_kernels, conv_strides)
callback = ModelCheckpoint(
monitor='val/f1',
mode='max',
save_top_k=5,
filename=CNN.__name__
)
trainer = pl.Trainer(
gpus=1,
max_epochs=150,
logger=logger,
callbacks=[callback]
)
trainer.fit(model, train_loader, val_loader)
# prediction = trainer.predict(model, test_loader)
# torch.save(prediction, 'sub.pt')
logger.finalize('ok')
if __name__ == '__main__':
train()
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