dev(narugo): simply refactor before running (e6186eb5) · Commits · git-mirror / Imgutils

zoo/monochrome/train_.py

+20 −11

Original line number	Diff line number	Diff line
		@@ -6,9 +6,9 @@ from typing import Optional, Type
		import torch
		from ditk import logging
		from torch import nn
		from torch.optim import lr_scheduler
		from torch.utils.data import DataLoader
		from torch.utils.tensorboard import SummaryWriter
		from torch.optim import lr_scheduler
		from tqdm.auto import tqdm

		from .alexnet import MonochromeAlexNet
		@@ -16,7 +16,7 @@ from .dataset import MonochromeDataset
		from .resnet import ResNet18, ResNet34, ResNet50, ResNet101, ResNet152
		from .transformer import SigTransformer
		from ..base import _TRAIN_DIR as _GLOBAL_TRAIN_DIR
		from ..utils import LRTyping, get_init_lr, get_dynamic_lr_scheduler
		from ..utils import LRTyping, get_init_lr

		_TRAIN_DIR = os.path.join(_GLOBAL_TRAIN_DIR, 'monochrome')
		_LOG_DIR = os.path.join(_TRAIN_DIR, 'logs')
		@@ -71,7 +71,7 @@ def _ckpt_epoch(filename: Optional[str]) -> Optional[int]:

		def train(dataset_dir: str, session_name: Optional[str] = None, from_ckpt: Optional[str] = None,
		train_ratio: float = 0.8, batch_size: int = 4, feature_bins: int = 256, fc: Optional[int] = 100,
		max_epochs: int = 500, learning_rate: LRTyping = 0.001, num_workers=8,
		max_epochs: int = 500, learning_rate: LRTyping = 0.001, num_workers: Optional[int] = None,
		save_per_epoch: int = 10, model_name: str = 'alexnet'):
		session_name = session_name or model_name
		_log_dir = os.path.join(_LOG_DIR, session_name)
		@@ -91,6 +91,7 @@ def train(dataset_dir: str, session_name: Optional[str] = None, from_ckpt: Optio
		test_size = dataset_size - train_size

		# 使用 random_split 函数拆分数据集
		num_workers = num_workers or os.cpu_count()
		train_dataset, test_dataset = torch.utils.data.random_split(full_dataset, [train_size, test_size])
		train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
		test_dataloader = DataLoader(test_dataset, batch_size=batch_size, num_workers=num_workers)
		@@ -107,17 +108,24 @@ def train(dataset_dir: str, session_name: Optional[str] = None, from_ckpt: Optio
		else:
		logging.info(f'No checkpoint found, new model will be used.')

		# Try use cude
		# Try use cuda
		if torch.cuda.is_available():
		model = model.cuda()

		loss_fn = nn.CrossEntropyLoss()
		initial_lr = get_init_lr(learning_rate)
		optimizer = torch.optim.AdamW([{'params': model.parameters(), 'initial_lr': initial_lr}], lr=initial_lr, weight_decay=1e-2)
		optimizer = torch.optim.AdamW(
		[{'params': model.parameters(), 'initial_lr': initial_lr}],
		lr=initial_lr, weight_decay=1e-2
		)
		# scheduler = get_dynamic_lr_scheduler(optimizer, lr=learning_rate, last_epoch=previous_epoch)
		scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=learning_rate,
		steps_per_epoch=len(train_dataloader), epochs=max_epochs,
		pct_start=0.15)
		scheduler = lr_scheduler.OneCycleLR(
		optimizer,
		max_lr=learning_rate,
		steps_per_epoch=len(train_dataloader),
		epochs=max_epochs,
		pct_start=0.15
		)

		for epoch in range(previous_epoch + 1, max_epochs + 1):
		running_loss = 0.0
		@@ -136,7 +144,8 @@ def train(dataset_dir: str, session_name: Optional[str] = None, from_ckpt: Optio
		scheduler.step()

		epoch_loss = running_loss / len(train_dataset)
		logging.info(f'Epoch [{epoch}/{max_epochs+1}] loss: {epoch_loss:.4f}, with learning rate: {scheduler.get_last_lr()[0]:.6f}')
		logging.info(f'Epoch [{epoch}/{max_epochs + 1}] loss: {epoch_loss:.4f}, '
		f'with learning rate: {scheduler.get_last_lr()[0]:.6f}')
		# scheduler.step()
		writer.add_scalar('train/loss', epoch_loss, epoch)

zoo/monochrome/transformer.py

+7 −7

File changed.

Contains only whitespace changes.