DataModule представляет собой класс, который создает загрузчики данных для обучения, валидации и тестирования модели.
config: Экземпляр класса TrainingConfig, который содержит конфигурационные параметры для обучения модели, такие как размер изображения, размер пакета для обучения, количество рабочих потоков и т. д.Пример данных
prepare_data()Загружает данные для обучения, валидации и тестирования. Параметры для загрузки данных (например, путь к данным) берутся из конфигурации.
setup()Разделяет данные на обучающий, валидационный и тестовый наборы. Здесь также применяются преобразования данных, такие как изменение размера, горизонтальное отражение и нормализация.
train_dataloader(), val_dataloader(), test_dataloader()Создают загрузчики данных для обучения, валидации и тестирования соответственно.
Класс UNet2DLightning представляет собой реализацию многоканальной UNet модели в PyTorch Lightning для задачи обработки изображений.
config: Экземпляр класса TrainingConfig, который содержит конфигурационные параметры для обучения модели, такие как размер изображения, размер пакета для обучения, количество рабочих потоков и т. д.Пример данных
__init__()Инициализирует объект класса. Создает экземпляры необходимых компонентов модели, такие как слои, функции потерь и расписывает гиперпараметры.
forward()Прямой проход модели. Принимает на вход изображение x, временные шаги timesteps и метки класса class_labels. Если pred=True, возвращает сэмплы изображений с добавленным шумом, иначе возвращает предсказания модели.
on_train_start(), on_training_epoch_end(), validation_step(), on_validation_epoch_end(), test_step()Методы, используемые для обучения, валидации и тестирования модели. Реализуют соответствующие этапы обучения и вычисления функции потерь.
predict_step()Метод для предсказания модели. Принимает на вход метки класса class_labels и возвращает сгенерированные изображения.
configure_optimizers()Конфигурирует оптимизатор и планировщик шага обучения для модели.
Этот тест проверяет, что при инициализации объекта DataModule с заданной конфигурацией объект конфигурации сохраняется внутри модуля.
Данный тест проверяет, что после вызова метода prepare_data() объект dataset в DataModule не является пустым.
Этот тест проверяет, что при передаче недопустимого имени датасета в метод prepare_data(), вызывается исключение ImportError.
@given(valid_training_config())
def test_setup(config):
data_module = DataModule(config)
data_module.prepare_data()
data_module.setup()
assert data_module.train_dataset is not None
assert data_module.test_dataset is not None
assert data_module.val_dataset is not None
assert isinstance(
data_module.train_dataloader(), torch.utils.data.dataloader.DataLoader
)
assert isinstance(
data_module.val_dataloader(), torch.utils.data.dataloader.DataLoader
)
assert isinstance(
data_module.test_dataloader(), torch.utils.data.dataloader.DataLoader
)Этот тест проверяет, что после подготовки данных методом setup(), обучающий, тестовый и валидационный наборы данных не являются пустыми, и что загрузчики данных создаются правильно.
@given(model_config())
def test_model_training(config):
config.num_timesteps = 1
model = UNet2DLightning(config=config)
image = torch.randn(
(
config.pred_batch_size,
1,
config.image_size,
config.image_size,
),
device=model.device,
)
rand_label = torch.tensor(0)
batch = [image, rand_label]
train_loss = model.training_step(batch=batch, batch_idx=0)
assert train_loss.item() >= 0Этот тест проверяет, что функция потерь на этапе обучения неотрицательна.
@given(model_config())
def test_model_validation(config):
config.num_timesteps = 1
model = UNet2DLightning(config=config)
image = torch.randn(
(
config.pred_batch_size,
1,
config.image_size,
config.image_size,
),
device=model.device,
)
rand_label = torch.tensor(0)
batch = [image, rand_label]
val_loss = model.validation_step(batch=batch, batch_idx=0)
assert val_loss.item() >= 0Этот тест проверяет, что функция потерь на этапе валидации неотрицательна.
@given(model_config())
def test_model_testing(config):
config.num_timesteps = 1
model = UNet2DLightning(config=config)
image = torch.randn(
(
config.pred_batch_size,
1,
config.image_size,
config.image_size,
),
device=model.device,
)
rand_label = torch.tensor(0)
batch = [image, rand_label]
test_loss = model.test_step(batch=batch, batch_idx=0)
assert test_loss.item() >= 0Этот тест проверяет, что функция потерь на этапе тестирования неотрицательна.
@given(model_config())
def test_model_prediction(config):
config.num_timesteps = 1
model = UNet2DLightning(config=config)
rand_label = torch.tensor(0)
with torch.no_grad():
pred_image = model.predict_step(batch=rand_label, batch_idx=0)
assert pred_image.shape == (
config.pred_batch_size,
config.image_size,
config.image_size,
1,
)Этот тест проверяет, что размер предсказанного изображения соответствует ожидаемому размеру.
При проверке тестов была обнаружена ошибка в тесте test_prepare_data_with_invalid_dataset, который проверяет обработку исключения ImportError при передаче недопустимого имени датасета в метод prepare_data().
Для исправления этой ошибки был использован следующий код:
Этот код позволяет проверить, существует ли модуль датасета с указанным именем, и если нет, то выбрасывает исключение ImportError с информацией о недоступности датасета.
После применения этого исправления, тест test_prepare_data_with_invalid_dataset корректно отработал и проверил обработку недопустимого имени датасета.