81 lines
2.9 KiB
Python
81 lines
2.9 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
|
sys.path.append(MODULE_ROOT)
|
|
|
|
from pytorch_wrapper.ptn_conv_model import PtnConvModel
|
|
from pytorch_wrapper.pto_adam import PtoAdam
|
|
from pytorch_wrapper.ptv_dataset import PtvDataset
|
|
from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor
|
|
from pytorch_wrapper.pt_data_loader import PtDataLoader
|
|
from pytorch_wrapper.pt_train_model import PtTrainModel
|
|
from pytorch_wrapper.pt_save_model import PtSaveModel
|
|
from pytorch_wrapper.pto_lr_scheduler_step import PtoLrSchedulerStep
|
|
from pytorch_wrapper.ptn_cross_entropy_loss import PtnCrossEntropyLoss
|
|
|
|
|
|
class TestPtTrainModel(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test instance."""
|
|
self.loss_fun = PtnCrossEntropyLoss().f(reduction="mean")[0]
|
|
compose_node = PtvTransformsToTensor()
|
|
transform = compose_node.f()[0]
|
|
dataset_node = PtvDataset()
|
|
dataset = dataset_node.f("FashionMNIST",
|
|
True, # download
|
|
"", # root
|
|
'{"train": True}',
|
|
transform=transform)[0]
|
|
data_loader_node = PtDataLoader()
|
|
self.train_loader = data_loader_node.f(
|
|
dataset=dataset,
|
|
batch_size=32,
|
|
shuffle=True,
|
|
parameters='{"num_workers":1}'
|
|
)[0]
|
|
|
|
self.model_node = PtnConvModel()
|
|
self.model = self.model_node.f(
|
|
input_dim="(1, 28, 28)",
|
|
penultimate_dim=0,
|
|
output_dim=10,
|
|
channel_list="[32,64,128,256]",
|
|
kernel_size_list="[3,3,3,3]",
|
|
padding_list="[1,1,1,1]",
|
|
downsample_list="[True,True,True,True]"
|
|
)[0]
|
|
|
|
self.optimizer = PtoAdam().f(self.model, 0.001, 0.9, 0.999)[0]
|
|
|
|
self.scheduler = PtoLrSchedulerStep().f(self.optimizer,
|
|
3, # steps
|
|
0.1 # gamma
|
|
)[0]
|
|
|
|
self.trainer = PtTrainModel()
|
|
self.save_model = PtSaveModel()
|
|
|
|
def test_1(self):
|
|
epochs = 4
|
|
trained_model = self.trainer.f(
|
|
self.model,
|
|
self.train_loader,
|
|
self.optimizer,
|
|
self.loss_fun,
|
|
epochs, # epochs
|
|
use_gpu=True,
|
|
early_stopping=False,
|
|
early_stopping_rounds=1,
|
|
scheduler=self.scheduler,
|
|
output_best_val_model=False)[0]
|
|
self.save_model.f(trained_model, f"fashion_mnist_{epochs}_epochs_conv.pt")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|