142 lines
3.4 KiB
Python
142 lines
3.4 KiB
Python
import os
|
|
import torch
|
|
import mlflow
|
|
import pytorch_lightning as pl
|
|
|
|
from torch import optim
|
|
from dotenv import load_dotenv
|
|
from torchvision.transforms import v2
|
|
from datasets import ImageNetDataModule
|
|
from train import ClassificationLightning
|
|
from custom_arhitecture import VGGCustom
|
|
|
|
from pytorch_lightning.loggers import MLFlowLogger
|
|
from pytorch_lightning.callbacks import LearningRateMonitor
|
|
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
|
|
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
|
|
|
|
|
|
load_dotenv()
|
|
|
|
IMAGENET_MEAN = [0.485, 0.456, 0.406]
|
|
IMAGENET_STD = [0.229, 0.224, 0.225]
|
|
|
|
dataset_root = "/media/daniel/developments/Datasets/ImageNet"
|
|
batch_size = 48
|
|
num_workers = 8
|
|
num_epoch = 90
|
|
learning_rate = 3e-4
|
|
|
|
train_transform = v2.Compose([
|
|
v2.RandomResizedCrop(
|
|
size=(224, 224),
|
|
scale=(0.08, 1.0),
|
|
ratio=(3 / 4, 4 / 3),
|
|
antialias=True
|
|
),
|
|
v2.RandomHorizontalFlip(p=0.5),
|
|
v2.ToDtype(torch.float32, scale=True),
|
|
v2.Normalize(
|
|
mean=IMAGENET_MEAN,
|
|
std=IMAGENET_STD
|
|
)
|
|
])
|
|
|
|
eval_transform = v2.Compose([
|
|
v2.Resize(256, antialias=True),
|
|
v2.CenterCrop(224),
|
|
v2.ToDtype(torch.float32, scale=True),
|
|
v2.Normalize(
|
|
mean=IMAGENET_MEAN,
|
|
std=IMAGENET_STD
|
|
)
|
|
])
|
|
|
|
|
|
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_kwargs = {
|
|
"lr": 3e-4,
|
|
"weight_decay": 1e-4
|
|
}
|
|
|
|
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"
|
|
)
|
|
]
|
|
|
|
# Системные метрики
|
|
mlflow.enable_system_metrics_logging()
|
|
mlflow.set_system_metrics_sampling_interval(5)
|
|
|
|
logger = MLFlowLogger(
|
|
experiment_name=os.getenv("EXPERIMENT_NAME"),
|
|
run_name="VGG16Custom",
|
|
tracking_uri=os.getenv("TRACKING_URI"),
|
|
log_model=True,
|
|
)
|
|
|
|
|
|
vgg = VGGCustom(architecture="VGG16",
|
|
num_channels=3,
|
|
num_classes=1000,
|
|
batch_normalization=False,
|
|
dropout=.5)
|
|
|
|
# Необходимо чтобы отрпаботал LazyLinear и принял размерности!
|
|
# with torch.no_grad():
|
|
# alexnet(torch.zeros(1, 3, 224, 224))
|
|
|
|
model = ClassificationLightning(model=vgg,
|
|
num_classes=1000,
|
|
optimizer_class=optim.Adam,
|
|
optimizer_kwargs=optimizer_Adam_kwargs,
|
|
scheduler_class=optim.lr_scheduler.ReduceLROnPlateau,
|
|
scheduler_kwargs=scheduler_ReduceLR_kwargs,
|
|
scheduler_monitor="val/loss")
|
|
|
|
trainer = pl.Trainer(
|
|
max_epochs=num_epoch,
|
|
accelerator='gpu',
|
|
devices=1,
|
|
precision="16-mixed",
|
|
logger=logger,
|
|
callbacks= callbacks
|
|
)
|
|
|
|
trainer.fit(model=model,
|
|
datamodule=imagenet_dm) |