xautodl/tests/test_torch_gpu_bugs.py

#####################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
#####################################################
# pytest ./tests/test_torch_gpu_bugs.py::test_create
#
# CUDA_VISIBLE_DEVICES="" pytest ./tests/test_torch_gpu_bugs.py::test_load
#####################################################
import os, sys, time, torch
import pickle
import tempfile
from pathlib import Path

root_dir = (Path(__file__).parent / ".." / "..").resolve()

from xautodl.trade_models.quant_transformer import QuantTransformer


def test_create():
    """Test the basic quant-model."""
    if not torch.cuda.is_available():
        return
    quant_model = QuantTransformer(GPU=0)
    temp_dir = root_dir / "tests" / ".pytest_cache"
    temp_dir.mkdir(parents=True, exist_ok=True)
    temp_file = temp_dir / "quant-model.pkl"
    with temp_file.open("wb") as f:
        # quant_model.to(None)
        quant_model.to("cpu")
        # del quant_model.model
        # del quant_model.train_optimizer
        pickle.dump(quant_model, f)
    print("save into {:}".format(temp_file))


def test_load():
    temp_file = root_dir / "tests" / ".pytest_cache" / "quant-model.pkl"
    with temp_file.open("rb") as f:
        model = pickle.load(f)
        print(model.model)
        print(model.train_optimizer)