Добавлен AlexNet

This commit is contained in:
Даниил Грабарь
2026-09-18 18:08:55 +10:00
parent d689368d75
commit 53b9c6d854
3 changed files with 265 additions and 62 deletions
+207 -48
View File
@@ -3,15 +3,21 @@
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, random_split
from torch.utils.data import DataLoader, Dataset
# Логирование данных
from pytorch_lightning.loggers import MLFlowLogger
# автоматическое отслеживание lr
from pytorch_lightning.callbacks import LearningRateMonitor
# ранние остановки
@@ -20,10 +26,11 @@ from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
# %% Параметры
dataset_root = "./Datasets/ImageNet"
batch_size = 64
dataset_root = "/media/daniel/developments/Datasets/ImageNet"
batch_size = 128
num_workers = 8
num_epoch = 2
num_epoch = 90
learning_rate = 3e-4
# %% Трансформации
@@ -59,61 +66,135 @@ path = kagglehub.competition_download('imagenet-object-localization-challenge',
print("Path to competition files:", path)
# %% Датасет
full_train_dataset = datasets.ImageNet(
root=dataset_root,
split="train",
transform=train_transform
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)
)
train_size = int(0.9 * len(full_train_dataset))
val_size = len(full_train_dataset) - train_size
def __len__(self):
return len(self.samples)
train_dataset, val_dataset = random_split(
full_train_dataset,
[train_size, val_size],
generator=torch.Generator().manual_seed(42)
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
)
val_base_dataset = datasets.ImageNet(
root=dataset_root,
split="train",
transform=eval_transform
)
val_dataset.dataset = val_base_dataset
test_dataset = datasets.ImageNet(
root=dataset_root,
split="val",
transform=eval_transform
)
train_loader = DataLoader(
dataset=train_dataset,
batch_size=batch_size,
def train_dataloader(self):
return DataLoader(
dataset=self.train_dataset,
batch_size=self.batch_size,
shuffle=True,
num_workers=num_workers,
num_workers=self.num_workers,
pin_memory=True,
persistent_workers=True
persistent_workers=self.num_workers > 0,
prefetch_factor=4 if self.num_workers > 0 else None
)
val_loader = DataLoader(
dataset=val_dataset,
batch_size=batch_size,
def val_dataloader(self):
return DataLoader(
dataset=self.val_dataset,
batch_size=self.batch_size,
shuffle=False,
num_workers=num_workers,
num_workers=self.num_workers,
pin_memory=True,
persistent_workers=True
persistent_workers=self.num_workers > 0,
prefetch_factor=4 if self.num_workers > 0 else None
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=batch_size,
def test_dataloader(self):
return DataLoader(
dataset=self.val_dataset,
batch_size=self.batch_size,
shuffle=False,
num_workers=num_workers,
num_workers=self.num_workers,
pin_memory=True,
persistent_workers=True
persistent_workers=self.num_workers > 0
)
# %% AlexNet
@@ -163,7 +244,7 @@ class AlexNetCustom(nn.Module):
def forward(self, image):
x = self.feature(image)
# x = self.avgpool(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.mlp(x)
@@ -241,9 +322,6 @@ class ClassificationLightning(pl.LightningModule):
stage,
metrics):
x, y = sample
x = x.to(torch.float32)
y = y.to(torch.int64)
logits = self(x)
loss = self.criterion(logits, y)
@@ -327,3 +405,84 @@ class ClassificationLightning(pl.LightningModule):
"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)
+34 -3
View File
@@ -11,6 +11,7 @@ dependencies:
- aiosignal=1.4.0=py312h06a4308_0
- alabaster=1.0.0=py312h06a4308_0
- alsa-lib=1.2.15.3=h47b2149_0
- annotated-types=0.7.0=py312h06a4308_0
- anyio=4.12.1=py312h06a4308_0
- aom=3.13.2=h664349e_0
- argon2-cffi=25.1.0=py312h06a4308_0
@@ -42,6 +43,7 @@ dependencies:
- charset-normalizer=3.4.7=py312h06a4308_0
- click=8.5.0=py312h06a4308_0
- cloudpickle=3.1.2=py312h06a4308_0
- coloredlogs=15.0.1=py312h06a4308_3
- comm=0.2.3=py312h06a4308_0
- cookiecutter=2.7.1=py312h06a4308_0
- cpython=3.12.14=py312hd8ed1ab_3
@@ -82,6 +84,8 @@ dependencies:
- gettext=0.25.1=h92eb808_0
- gettext-tools=0.25.1=h6a67909_0
- giflib=6.1.3=h621bd71_1
- gitdb=4.0.12=py312h06a4308_0
- gitpython=3.1.59=py312h06a4308_0
- gmp=6.3.0=hfd2156b_3
- gmpy2=2.3.1=py312h04c1168_0
- graphite2=1.3.15=h9ba177b_0
@@ -89,6 +93,7 @@ dependencies:
- harfbuzz=12.3.0=h79d275a_1
- httpcore=1.0.9=py312h06a4308_0
- httpx=0.28.1=py312h06a4308_1
- humanfriendly=10.0=py312h06a4308_2
- icu=73.1=h6a678d5_0
- idna=3.18=py312h06a4308_0
- imagesize=1.4.1=py312h06a4308_0
@@ -98,6 +103,7 @@ dependencies:
- ipykernel=6.31.0=py312h7040dfc_0
- ipython=9.15.0=py312h06a4308_0
- ipython_pygments_lexers=1.1.1=py312h06a4308_0
- ipywidgets=8.1.9=py312h06a4308_0
- isort=8.0.1=py312h06a4308_1
- jansson=2.15.0=hbcba0ee_0
- jaraco.classes=3.4.0=py312h06a4308_0
@@ -121,6 +127,7 @@ dependencies:
- jupyterlab=4.6.2=py312h06a4308_1
- jupyterlab_pygments=0.3.0=py312h06a4308_0
- jupyterlab_server=2.28.0=py312h06a4308_1
- jupyterlab_widgets=3.0.17=pyhd3eb1b0_0
- kagglehub=1.0.2=py312h06a4308_0
- kagglesdk=0.1.35=py312h06a4308_0
- keyring=25.7.0=py312h06a4308_0
@@ -219,6 +226,7 @@ dependencies:
- mesalib=25.1.5=h3583ad3_5
- minizip=4.2.1=h7e49174_0
- mistune=3.3.3=py312h06a4308_0
- ml_dtypes=0.6.0=py312h86c3e14_0
- more-itertools=11.0.2=py312h06a4308_0
- mpc=1.4.0=ha2cb11d_1
- mpfr=4.2.2=ha2cb11d_1
@@ -248,11 +256,16 @@ dependencies:
- numpy=2.5.2=py312h35deafb_0
- numpy-base=2.5.2=py312h4bc27c9_0
- numpydoc=1.10.0=py312h06a4308_0
- onnx=1.22.0=py312hc50a595_0
- onnx-ir=0.2.1=py312h06a4308_0
- onnxruntime=1.24.4=py312h0197eaa_2_cpu
- onnxscript=0.7.1=py312h06a4308_0
- openjpeg=2.5.4=h4e0627c_1
- openldap=2.6.12=h007892f_1
- openssl=3.6.4=h781a0a9_0
- optree=0.20.0=py312ha5105d2_1
- packaging=26.3=py312h06a4308_0
- pandas=3.0.5=py312h86c3e14_0
- pandoc=3.8=h06a4308_0
- pandocfilters=1.5.1=py312h06a4308_0
- parso=0.8.7=py312h06a4308_0
@@ -280,6 +293,8 @@ dependencies:
- pycares=4.10.0=py312ha137689_0
- pycodestyle=2.12.1=py312h06a4308_0
- pycparser=3.0=py312h06a4308_0
- pydantic=2.13.4=py312h06a4308_0
- pydantic-core=2.46.4=py312h6e1b9ff_0
- pydocstyle=6.3.0=py312h06a4308_1
- pyflakes=3.2.0=py312h06a4308_0
- pygithub=2.9.1=py312h06a4308_0
@@ -296,6 +311,7 @@ dependencies:
- python=3.12.14=hc5f7cf0_0
- python-dateutil=2.9.0post0=py312h06a4308_2
- python-fastjsonschema=2.22.1=py312h06a4308_0
- python-flatbuffers=25.9.23=py312h8034063_0
- python-json-logger=4.1.0=py312h06a4308_0
- python-lsp-black=2.0.0=py312h06a4308_1
- python-lsp-jsonrpc=1.1.2=pyhd3eb1b0_0
@@ -318,6 +334,7 @@ dependencies:
- qtbase=6.11.0=h9201cad_0
- qtconsole=5.7.2=py312h06a4308_0
- qtdeclarative=6.11.0=h89ee561_1
- qtimageformats=6.11.0=h00148ef_1
- qtpy=2.4.3=py312h06a4308_0
- qtsvg=6.11.0=hbecefc2_1
- qttools=6.11.0=h0af0a4d_1
@@ -338,12 +355,15 @@ dependencies:
- rpds-py=0.28.0=py312h498d7c9_0
- rtree=1.4.1=py312hd966e25_1
- ruff=0.15.1=py312h6e1b9ff_0
- scipy=1.18.1=py312h35deafb_0
- secretstorage=3.5.0=py312h3e8c6aa_0
- send2trash=1.8.3=py312h06a4308_0
- sentry-sdk=2.67.1=py312h06a4308_0
- setuptools=80.10.2=py312h06a4308_0
- sip=6.15.1=py312he6c97cb_0
- six=1.17.0=py312h06a4308_0
- sleef=3.9.0=ha0421bc_0
- smmap=5.0.2=py312h06a4308_0
- snappy=1.2.2=h4bcf44c_1
- snowballstemmer=3.1.1=py312h06a4308_0
- sortedcontainers=2.4.0=pyhd3eb1b0_0
@@ -377,11 +397,12 @@ dependencies:
- traitlets=5.15.0=py312h06a4308_0
- triton=3.7.0=cuda129py312h338dc61_1
- typing-extensions=4.16.0=py312h06a4308_0
- typing-inspection=0.4.3=py312h06a4308_0
- typing_extensions=4.16.0=py312h06a4308_0
- tzdata=2026c=he532380_0
- ujson=5.13.0=py312h7354ed3_0
- urllib3=2.7.0=py312h06a4308_0
- uvloop=0.22.1=py312h939ff40_1
- wandb=0.27.2=py312h03481c5_1
- watchdog=6.0.0=py312h06a4308_0
- wayland=1.24.0=hdac8c69_0
- wcwidth=0.8.2=py312h06a4308_0
@@ -389,9 +410,10 @@ dependencies:
- websocket-client=1.8.0=py312h06a4308_0
- whatthepatch=1.0.7=py312h06a4308_0
- wheel=0.47.0=py312h06a4308_0
- widgetsnbextension=4.0.16=py312h06a4308_0
- wurlitzer=3.1.1=py312h06a4308_0
- xcb-util=0.4.1=h5eee18b_2
- xcb-util-cursor=0.1.5=h5eee18b_0
- xcb-util-cursor=0.1.6=hb03c661_0
- xcb-util-image=0.4.0=h5eee18b_2
- xcb-util-keysyms=0.4.1=h5eee18b_0
- xcb-util-renderutil=0.3.10=h5eee18b_0
@@ -422,6 +444,15 @@ dependencies:
- zlib=1.3.2=h47b2149_0
- zstd=1.5.7=h11fc155_0
- pip:
- pyqt6-qt6==6.11.2
- backoff==2.2.1
- blake3==1.0.9
- docker==7.2.0
- lightning-sdk==2026.9.15
- linkify-it-py==2.2.0
- litlogger==2026.8.28
- mdit-py-plugins==0.6.1
- rich-click==1.9.9
- split-folders==0.6.1
- textual==8.2.8
- urllib3==2.5.0
prefix: /home/daniel/anaconda3/envs/CV
+13
View File
@@ -10,4 +10,17 @@ dependencies:
- beartype
- torchvision
- kagglehub
- xcb-util-cursor
- pyqt
- pyqt6-sip
- pyqtwebengine
- qtbase
- qtimageformats
- onnx
- onnxruntime
- onnxscript
- scipy
- ipywidgets
- pandas
- wandb
prefix: /home/daniel/anaconda3/envs/CV