e2e_gdn/e2e_gdn/test_ms-all.py

import torch
import mindspore as ms
import mindspore.nn as nn
from mindspore import save_checkpoint, Tensor
from mindspore import load_checkpoint, load_param_into_net
from model import ImageCompressor
from PIL import Image
from pytorch_msssim import ms_ssim
import os
import glob
import time
import pandas as pd
import numpy as np
from mindspore import context

ms.set_context(device_target="GPU")
context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU") #GRAPH_MODE(静态图模式) PYNATIVE_MODE(动态图模式)
context.set_context(save_graphs=False)
context.set_context(device_id=int(os.getenv('DEVICE_ID', '0')))
if ms.get_context("device_target") == "GPU":
    context.set_context(enable_graph_kernel=False) #开启图算融合以优化网络执行性能，常用于GPU，动静态模式应该都可以
if ms.get_context("mode") == ms.PYNATIVE_MODE:
    ms.set_context(mempool_block_size="1GB") #设置设备内存池的块大小，实际使用的内存池块大小是设备的可用内存和 mempool_block_size 值中的最小值
ms.reset_auto_parallel_context()
ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.STAND_ALONE, gradients_mean=True, device_num=1)
def ms_compute_psnr(a, b):
    ms_psnr = ms.nn.PSNR(1.0)
    return ms_psnr(a, b)

def ms_compute_msssim(a, b):
    m_ssim = ms.nn.MSSSIM(max_val=1.0, power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333), filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03)
    return m_ssim(a, b)

if not os.path.exists('test_output'):
    os.makedirs('test_output')
model_list = ['1500_iter_19762.pth.ckpt','400_iter_19762.pth.ckpt','800_iter_19762.pth.ckpt','3000_iter_19844.pth.ckpt']
imglist = glob.glob("/userhome/dataset/kodak/test/*.png")
model = ImageCompressor()
for model_name in model_list:
    out_dir = 'test_output/lambda_' + model_name.split('_')[0]
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    param_dict = load_checkpoint(model_name)
    param_not_load = load_param_into_net(model, param_dict)
    print("param_not_load: ", param_not_load) #打印网络中没有被加载的参数，正常应该为空
    model.set_train(False)
    path_all = np.array([])
    PSNR_all = np.array([])
    bpp_all = np.array([])
    MSSSIM_all = np.array([])
    run_time = np.array([])
    for i, oneimg in enumerate(imglist):
        start_time = time.time()
        print('enc '+oneimg.split('/')[-1])
        path_all = np.concatenate((path_all, [oneimg.split('/')[-1]]), axis=0)
        img = Image.open(oneimg).convert('RGB')
        totensor = ms.dataset.vision.ToTensor()
        data = ms.Tensor(totensor(img)).unsqueeze(0)
        print(data.shape)
        out = model(data)
        afterdec_time = time.time()
        run_time = np.concatenate((run_time, [afterdec_time-start_time]), axis=0)
        test_psnr = ms_compute_psnr(data, out[0]).asnumpy()
        PSNR_all = np.concatenate((PSNR_all, test_psnr), axis=0)
        test_bpp = out[2].asnumpy()
        bpp_all = np.concatenate((bpp_all, [test_bpp]), axis=0)
        ms_ssim = ms_compute_msssim(data, out[0]).asnumpy()
        MSSSIM_all = np.concatenate((MSSSIM_all, ms_ssim), axis=0)
        # if i>1:break
    path_all = np.concatenate((path_all, ['all']), axis=0)
    path_all = path_all.reshape(-1, 1)
    PSNR_all = np.concatenate((PSNR_all, [np.mean(PSNR_all)]), axis=0)
    PSNR_all = PSNR_all.reshape(-1, 1)
    bpp_all = np.concatenate((bpp_all, [np.mean(bpp_all)]), axis=0)
    bpp_all = bpp_all.reshape(-1, 1)
    MSSSIM_all = np.concatenate((MSSSIM_all, [np.mean(MSSSIM_all)]), axis=0)
    MSSSIM_all = MSSSIM_all.reshape(-1, 1)
    run_time = np.concatenate((run_time, [np.sum(run_time)]), axis=0)
    run_time = run_time.reshape(-1, 1)
    all_results = np.concatenate((path_all, bpp_all, PSNR_all, MSSSIM_all,run_time), axis=1)
    results2 = pd.DataFrame(columns=['imgname', 'bpp', 'PSNR', 'MSSSIM','run_time'], data=all_results)
    results2.to_csv(f'test_output/e2e_gdn_mindspore_lambda'+model_name.split('_')[0]+'.csv',index=False)