489 lines
14 KiB
Python
489 lines
14 KiB
Python
# %% Импорт библиотек
|
|
|
|
import time
|
|
import torch
|
|
import kagglehub
|
|
import pandas as pd
|
|
import torch.nn as nn
|
|
import pytorch_lightning as pl
|
|
|
|
from PIL import Image
|
|
from torch import optim
|
|
from pathlib import Path
|
|
from torchmetrics import MetricCollection
|
|
from typing import Optional, Callable, Any
|
|
from torchvision import datasets, transforms
|
|
from torchmetrics.classification import Accuracy
|
|
from torch.utils.data import DataLoader, Dataset
|
|
|
|
# Логирование данных
|
|
from pytorch_lightning.loggers import MLFlowLogger
|
|
# автоматическое отслеживание lr
|
|
from pytorch_lightning.callbacks import LearningRateMonitor
|
|
# ранние остановки
|
|
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
|
|
# сохранение весов
|
|
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
|
|
# %% Параметры
|
|
|
|
dataset_root = "/media/daniel/developments/Datasets/ImageNet"
|
|
batch_size = 128
|
|
num_workers = 8
|
|
num_epoch = 90
|
|
learning_rate = 3e-4
|
|
|
|
# %% Трансформации
|
|
|
|
train_transform = transforms.Compose([
|
|
transforms.RandomResizedCrop(224),
|
|
transforms.RandomHorizontalFlip(),
|
|
|
|
transforms.ToTensor(),
|
|
|
|
transforms.Normalize(
|
|
mean=[0.485, 0.456, 0.406],
|
|
std=[0.229, 0.224, 0.225]
|
|
)
|
|
])
|
|
|
|
eval_transform = transforms.Compose([
|
|
transforms.Resize(256),
|
|
transforms.CenterCrop(224),
|
|
|
|
transforms.ToTensor(),
|
|
|
|
transforms.Normalize(
|
|
mean=[0.485, 0.456, 0.406],
|
|
std=[0.229, 0.224, 0.225]
|
|
)
|
|
])
|
|
|
|
# %% Скачивание ImageNet
|
|
|
|
path = kagglehub.competition_download('imagenet-object-localization-challenge',
|
|
output_dir=dataset_root,
|
|
force_download=False)
|
|
print("Path to competition files:", path)
|
|
# %% Датасет
|
|
|
|
TransformType = Optional[Callable[..., Any]]
|
|
|
|
class ImageNetValDataset(Dataset):
|
|
def __init__(
|
|
self,
|
|
image_dir,
|
|
labels_file,
|
|
class_to_idx,
|
|
classes,
|
|
transform=None
|
|
):
|
|
self.image_dir = Path(image_dir)
|
|
self.transform = transform
|
|
|
|
self.class_to_idx = class_to_idx
|
|
self.classes = classes
|
|
|
|
df = pd.read_csv(labels_file)
|
|
|
|
self.samples = []
|
|
|
|
for _, row in df.iterrows():
|
|
image_id = row["ImageId"]
|
|
synset = row["PredictionString"].split()[0]
|
|
label = self.class_to_idx[synset]
|
|
image_path = self.image_dir / f"{image_id}.JPEG"
|
|
self.samples.append(
|
|
(image_path, label)
|
|
)
|
|
|
|
def __len__(self):
|
|
return len(self.samples)
|
|
|
|
def __getitem__(self, idx):
|
|
image_path, label = self.samples[idx]
|
|
|
|
image = Image.open(image_path).convert("RGB")
|
|
|
|
if self.transform is not None:
|
|
image = self.transform(image)
|
|
|
|
return image, label
|
|
|
|
|
|
class ImageNetDataModule(pl.LightningDataModule):
|
|
def __init__(self,
|
|
dataset_root: Path,
|
|
train_transform: TransformType,
|
|
eval_transform: TransformType,
|
|
batch_size: int = 64,
|
|
num_workers:int = 8
|
|
):
|
|
super().__init__()
|
|
|
|
self.dataset_root = Path(dataset_root)
|
|
|
|
self.train_transform = train_transform
|
|
self.eval_transform = eval_transform
|
|
self.num_workers = num_workers
|
|
self.batch_size = batch_size
|
|
|
|
self.imagenet_root = Path(dataset_root) / "ILSVRC"
|
|
|
|
self.train_dir = self.imagenet_root / "Data" / "CLS-LOC" / "train"
|
|
self.val_dir = self.imagenet_root / "Data" / "CLS-LOC" / "val"
|
|
|
|
self.mapping_file = Path(dataset_root) / "LOC_synset_mapping.txt"
|
|
self.val_labels_file = Path(dataset_root) / "LOC_val_solution.csv"
|
|
|
|
self.train_dataset = None
|
|
self.val_dataset = None
|
|
|
|
self.class_to_idx = None
|
|
self.classes = None
|
|
|
|
def setup(self, stage=None):
|
|
synset_to_name = {}
|
|
with open(self.mapping_file, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
synset, class_name = line.strip().split(" ", 1)
|
|
synset_to_name[synset] = class_name
|
|
|
|
train_dataset = datasets.ImageFolder(root=self.train_dir,
|
|
transform=self.train_transform)
|
|
|
|
self.class_to_idx = train_dataset.class_to_idx.copy()
|
|
self.classes = [synset_to_name[synset]
|
|
for synset in train_dataset.classes]
|
|
train_dataset.classes = self.classes
|
|
self.train_dataset = train_dataset
|
|
|
|
self.val_dataset = ImageNetValDataset(
|
|
image_dir=self.val_dir,
|
|
labels_file=self.val_labels_file,
|
|
class_to_idx=self.class_to_idx,
|
|
classes=self.classes,
|
|
transform=self.eval_transform
|
|
)
|
|
|
|
def train_dataloader(self):
|
|
return DataLoader(
|
|
dataset=self.train_dataset,
|
|
batch_size=self.batch_size,
|
|
shuffle=True,
|
|
num_workers=self.num_workers,
|
|
pin_memory=True,
|
|
persistent_workers=self.num_workers > 0,
|
|
prefetch_factor=4 if self.num_workers > 0 else None
|
|
)
|
|
|
|
def val_dataloader(self):
|
|
return DataLoader(
|
|
dataset=self.val_dataset,
|
|
batch_size=self.batch_size,
|
|
shuffle=False,
|
|
num_workers=self.num_workers,
|
|
pin_memory=True,
|
|
persistent_workers=self.num_workers > 0,
|
|
prefetch_factor=4 if self.num_workers > 0 else None
|
|
)
|
|
|
|
def test_dataloader(self):
|
|
return DataLoader(
|
|
dataset=self.val_dataset,
|
|
batch_size=self.batch_size,
|
|
shuffle=False,
|
|
num_workers=self.num_workers,
|
|
pin_memory=True,
|
|
persistent_workers=self.num_workers > 0
|
|
)
|
|
|
|
# %% AlexNet
|
|
|
|
class AlexNetCustom(nn.Module):
|
|
def __init__(self,
|
|
in_channels: int = 3,
|
|
out_classes: int=1_000):
|
|
super().__init__()
|
|
|
|
self._out_classes = out_classes
|
|
|
|
self.feature = nn.Sequential(
|
|
nn.Conv2d(in_channels, 96, kernel_size=11, stride=4, padding=2),
|
|
nn.ReLU(inplace=True),
|
|
nn.MaxPool2d(kernel_size=3, stride=2),
|
|
|
|
nn.Conv2d(96, 256, kernel_size=5, padding=2),
|
|
nn.ReLU(inplace=True),
|
|
nn.MaxPool2d(kernel_size=3, stride=2),
|
|
|
|
nn.Conv2d(256, 384, kernel_size=3, padding=1),
|
|
nn.ReLU(inplace=True),
|
|
|
|
nn.Conv2d(384, 384, kernel_size=3, padding=1),
|
|
nn.ReLU(inplace=True),
|
|
|
|
nn.Conv2d(384, 256, kernel_size=3, padding=1),
|
|
nn.ReLU(inplace=True),
|
|
nn.MaxPool2d(kernel_size=3, stride=2)
|
|
)
|
|
|
|
# Добавлено в современной реализации
|
|
self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
|
|
|
|
self.mlp = nn.Sequential(
|
|
nn.Dropout(p=0.5),
|
|
nn.Linear(256*6*6, 4096),
|
|
nn.ReLU(inplace=True),
|
|
|
|
nn.Dropout(p=0.5),
|
|
nn.Linear(4096, 4096),
|
|
nn.ReLU(inplace=True),
|
|
|
|
nn.Linear(4096, self._out_classes)
|
|
)
|
|
|
|
def forward(self, image):
|
|
x = self.feature(image)
|
|
x = self.avgpool(x)
|
|
x = torch.flatten(x, 1)
|
|
x = self.mlp(x)
|
|
|
|
return x
|
|
|
|
# %% Классификатор обучатор
|
|
|
|
class ClassificationLightning(pl.LightningModule):
|
|
def __init__(self,
|
|
model: nn.Module,
|
|
num_classes: int,
|
|
optimizer_class,
|
|
optimizer_kwargs: dict,
|
|
scheduler_class=None,
|
|
scheduler_kwargs: dict | None = None,
|
|
scheduler_monitor: str = "val/loss"):
|
|
super().__init__()
|
|
|
|
|
|
scheduler_kwargs = scheduler_kwargs or {}
|
|
|
|
self.model = model
|
|
self.optimizer_class = optimizer_class
|
|
self.scheduler_class = scheduler_class
|
|
|
|
num_parameters = sum(
|
|
p.numel()
|
|
for p in model.parameters()
|
|
)
|
|
|
|
num_trainable_parameters = sum(
|
|
p.numel()
|
|
for p in model.parameters()
|
|
if p.requires_grad
|
|
)
|
|
|
|
self.save_hyperparameters({
|
|
"model": model.__class__.__name__,
|
|
"num_classes": num_classes,
|
|
|
|
"num_parameters": num_parameters,
|
|
"num_trainable_parameters": num_trainable_parameters,
|
|
|
|
"optimizer": optimizer_class.__name__,
|
|
"optimizer_kwargs": optimizer_kwargs,
|
|
|
|
"scheduler": (
|
|
scheduler_class.__name__
|
|
if scheduler_class is not None
|
|
else None
|
|
),
|
|
"scheduler_kwargs": scheduler_kwargs,
|
|
"scheduler_monitor": scheduler_monitor
|
|
})
|
|
|
|
|
|
self.criterion = nn.CrossEntropyLoss()
|
|
self.train_metrics = MetricCollection({
|
|
"acc_top1": Accuracy(task="multiclass",
|
|
num_classes=self.hparams.num_classes,
|
|
top_k=1),
|
|
"acc_top5": Accuracy(task="multiclass",
|
|
num_classes=self.hparams.num_classes,
|
|
top_k=5),
|
|
})
|
|
|
|
self.val_metrics = self.train_metrics.clone()
|
|
self.test_metrics = self.train_metrics.clone()
|
|
|
|
def forward(self, image):
|
|
return self.model(image)
|
|
|
|
def _shared_step(self,
|
|
sample,
|
|
stage,
|
|
metrics):
|
|
x, y = sample
|
|
|
|
logits = self(x)
|
|
loss = self.criterion(logits, y)
|
|
|
|
metrics.update(logits, y)
|
|
|
|
self.log(f"{stage}/loss",
|
|
loss,
|
|
prog_bar=True,
|
|
on_epoch=True,
|
|
on_step=False
|
|
)
|
|
|
|
return loss
|
|
|
|
def _log_metrics(self, stage, metrics):
|
|
computed = metrics.compute()
|
|
|
|
for name, value in computed.items():
|
|
self.log(f"{stage}/{name}", value, prog_bar=True)
|
|
|
|
metrics.reset()
|
|
|
|
def on_train_epoch_start(self):
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
|
|
self._epoch_start_time = time.perf_counter()
|
|
|
|
def training_step(self, batch, batch_idx):
|
|
return self._shared_step(batch, "train", self.train_metrics)
|
|
|
|
def on_train_epoch_end(self):
|
|
self._log_metrics("train", self.train_metrics)
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
|
|
epoch_time = time.perf_counter() - self._epoch_start_time
|
|
lr = self.optimizers().param_groups[0]["lr"]
|
|
|
|
self.log("train/lr", lr, prog_bar=True)
|
|
self.log("train/epoch_time", epoch_time)
|
|
|
|
def validation_step(self, batch, batch_idx):
|
|
return self._shared_step(batch, "val", self.val_metrics)
|
|
|
|
def on_validation_epoch_end(self):
|
|
self._log_metrics("val", self.val_metrics)
|
|
|
|
def test_step(self, batch, batch_idx):
|
|
return self._shared_step(batch, "test", self.test_metrics)
|
|
|
|
def on_test_epoch_end(self):
|
|
self._log_metrics("test", self.test_metrics)
|
|
|
|
def configure_optimizers(self):
|
|
optimizer = self.optimizer_class(
|
|
self.parameters(),
|
|
**self.hparams.optimizer_kwargs
|
|
)
|
|
|
|
if self.scheduler_class is None:
|
|
return optimizer
|
|
|
|
scheduler = self.scheduler_class(
|
|
optimizer,
|
|
**self.hparams.scheduler_kwargs
|
|
)
|
|
|
|
scheduler_config = {
|
|
"scheduler": scheduler,
|
|
"interval": "epoch",
|
|
"frequency": 1
|
|
}
|
|
|
|
if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):
|
|
scheduler_config["monitor"] = self.hparams.scheduler_monitor
|
|
|
|
return {
|
|
"optimizer": optimizer,
|
|
"lr_scheduler": scheduler_config
|
|
}
|
|
|
|
# %% Обучение
|
|
|
|
pl.seed_everything(42, workers=True)
|
|
torch.set_float32_matmul_precision("high")
|
|
|
|
imagenet_dm = ImageNetDataModule(
|
|
dataset_root=dataset_root,
|
|
train_transform=train_transform,
|
|
eval_transform=eval_transform,
|
|
batch_size=batch_size,
|
|
num_workers=num_workers
|
|
)
|
|
|
|
|
|
optimizer_Adam = optim.Adam
|
|
optimizer_Adam_kwargs = {
|
|
"lr": 3e-4,
|
|
"weight_decay": 1e-4
|
|
}
|
|
|
|
scheduler_ReduceLR = optim.lr_scheduler.ReduceLROnPlateau
|
|
scheduler_ReduceLR_kwargs={
|
|
"mode": "min",
|
|
"factor": 0.1,
|
|
"patience": 5
|
|
}
|
|
|
|
callbacks = [
|
|
ModelCheckpoint(
|
|
dirpath="models",
|
|
filename="epoch={epoch}",
|
|
auto_insert_metric_name=False,
|
|
save_last=True,
|
|
save_top_k=1,
|
|
monitor="val/acc_top1",
|
|
mode="max"
|
|
),
|
|
|
|
LearningRateMonitor(
|
|
logging_interval="epoch"
|
|
),
|
|
|
|
EarlyStopping(
|
|
monitor="val/loss",
|
|
min_delta=2e-4,
|
|
patience=10,
|
|
verbose=False,
|
|
mode="min"
|
|
)
|
|
]
|
|
|
|
logger = MLFlowLogger(
|
|
experiment_name="ImageNetMyAlgoritms",
|
|
run_name="AlexNetCustom",
|
|
tracking_uri="http://10.16.18.1:8055",
|
|
log_model=True,
|
|
)
|
|
|
|
alexnet = AlexNetCustom(in_channels=3, out_classes=1000)
|
|
model_adam = ClassificationLightning(
|
|
model=alexnet,
|
|
num_classes=1000,
|
|
optimizer_class=optimizer_Adam,
|
|
optimizer_kwargs=optimizer_Adam_kwargs,
|
|
scheduler_class=scheduler_ReduceLR,
|
|
scheduler_kwargs=scheduler_ReduceLR_kwargs,
|
|
scheduler_monitor="val/loss"
|
|
)
|
|
|
|
trainer = pl.Trainer(
|
|
max_epochs=num_epoch,
|
|
accelerator="gpu",
|
|
devices=1,
|
|
logger=logger,
|
|
# precision="16-mixed",
|
|
callbacks=callbacks
|
|
)
|
|
|
|
trainer.fit(model=model_adam,
|
|
datamodule=imagenet_dm)
|