def binary_search_helper(dimension, logger, model_name="EDSR", device="cuda"):
"""
Process random image and calculates processing time
Parameters
----------
dimension : int
random image dimension.
logger : logger
keep logs.
device : str, optional
GPU or CPU. The default is 'cuda'.
Returns
-------
total_time : float
EDSR processing time.
"""
print('Before loading model: ')
subprocess.run("gpustat", shell=True)
print()
total_time = 0
try:
model = None
if model_name == "EDSR":
model = md.load_edsr(device=device)
print('After loading model: ')
subprocess.run("gpustat", shell=True)
print()
elif model_name == "RRDB":
model = md.load_rrdb(device=device)
else:
raise Exception("Unknown model...")
model.eval()
input_image = ut.random_image(dimension)
if model_name == "RRDB":
input_image = input_image[:, 2:, :, :]
input_image = input_image.to(device)
with torch.no_grad():
start = time.time()
print('Before processing: ')
subprocess.run("gpustat", shell=True)
output_image = model(input_image)
print('After processing: ')
subprocess.run("gpustat", shell=True)
end = time.time()
total_time = end - start
ut.clear_cuda(input_image, output_image)
model.cpu()
del model
print('After model shifting and deleting: ')
subprocess.run("gpustat", shell=True)
except RuntimeError as err:
logger.error("Runtime error for dimension: {}x{}: " + err)
sys.exit(1)
return total_time
def main(img_path, model_name, patch_dimension):
"""
Driver for recursive forward chop. Takes an image path and model name to upsample the image.
"""
# Loading model and image
img = None
model = None
print("\nLoading model and image... \n")
if model_name in ["EDSR"]:
img = ut.load_image(img_path)
img = img.unsqueeze(0)
model = md.load_edsr(device="cuda")
else:
img = ut.npz_loader(img_path)
img = img.unsqueeze(0)
model = md.load_rrdb(device="cuda")
# Timers for saving stats
timer = [0, 0, 0, 0, 0]
print("Processing...")
# Shfiting input image to CUDA
total_time = ut.timer()
cpu2gpu_time = ut.timer()
img = img.to("cuda")
cpu2gpu_time = cpu2gpu_time.toc()
# Forward chopping and upsampling
output = forward_chop(
img, model, timer=timer, min_size=patch_dimension * patch_dimension
)
# Shifting output image to CPU
gpu2cpu_time = ut.timer()
output = output.to("cpu")
gpu2cpu_time = gpu2cpu_time.toc()
if model_name in ["EDSR"]:
output = output.int()
total_time = total_time.toc()
timer[0] = cpu2gpu_time
timer[-2] = gpu2cpu_time
timer[-1] = total_time
# Saving output
np.savez("results/recursive_outputx4.npz", output)
np.save("results/recursive_outputx4", output)
# Printing processing times
print("\nCPU 2 GPU time: ", timer[0])
print("\nUpsampling time: ", timer[1])
print("\nMerging time: ", timer[2])
print("\nGPU 2 CPU time", timer[3])
print("\nTotal execution time: ", timer[4])
def helper_upsampler_piterative_experiment(model_name, img_path,
patch_dimension):
"""
Driver function for running pytorch model inference
Parameters
----------
img_dimension : int
image one side dimension.
patch_dimension : int
patch size.
Returns
-------
None.
"""
# Loading model and image
if model_name in ['EDSR']:
model = md.load_edsr(device="cuda")
img = ut.load_image(img_path)
input_image = img.unsqueeze(0)
elif model_name in ["RRDB"]:
model = md.load_rrdb(device="cuda")
img = ut.npz_loader(img_path)
input_image = img.unsqueeze(0)
else:
print('Unknown model!')
return
b, c, h, w = input_image.shape
total_time = ut.timer()
out_tuple = forward_chop_iterative(
input_image,
shave=10,
min_size=patch_dimension * patch_dimension,
model=model,
device="cuda",
print_result=True,
)
model.cpu()
del model
output_image = out_tuple[0]
total_time = total_time.toc()
# =============================================================================
# for i in out_tuple[1:]:
# print(i)
# =============================================================================
print('Total executing time: ', total_time)
return output_image
def main(model_name, onnx_model_name):
device = ut.get_device_type()
model = None
if model_name == "RRDB":
model = md.load_rrdb(device)
dummy_input = ut.random_image(100, batch=False, channel=1).to(device)
elif model_name == "EDSR":
model = md.load_edsr(device)
dummy_input = ut.random_image(100).to(device)
model.eval()
torch.onnx.export(model,
dummy_input,
"inference_models/" + onnx_model_name,
verbose=False)
if __name__ == "__main__":
img_path = sys.argv[1] if len(sys.argv) > 1 else "test2.jpg"
dimension = int(sys.argv[2]) if len(sys.argv) > 2 else 32
shave = int(sys.argv[3]) if len(sys.argv) > 3 else 12
device = "cuda"
img = torchvision.io.read_image(img_path)
c, h, w = img.shape
img = img.reshape((1, c, h, w))
plt.imshow(img[0].permute((1, 2, 0)))
input_image = img.float()
model = md.load_edsr(device=device)
model.eval()
st = time.time()
out = forward_chop_iterative(input_image,
shave=shave,
min_size=dimension * dimension,
model=model)
et = time.time()
tt = et - st
print("Total forward chopping time: ", tt)
out = out.int()
print(out.shape)
# =============================================================================
# from PIL import Image
# im = Image.fromarray(np.array(out[0].permute(1,2,0)))
# im.save("4xinput.jpg")
def build_onnx_model(model_name, patch_size, onnx_model_name, device="cuda"):
"""
Builds ONNX model with a fixed input shape
Parameters
----------
model_name : str
Upsampler model name. (i.e. RRDB, EDSR).
patch_size : int
Input image dimension n of nxn.
onnx_model_name : str
output ONNX model name.
Returns
-------
None.
"""
# =============================================================================
# print('Before loading model: ')
# subprocess.run("gpustat", shell = True)
# print()
# total_time = 0
# model = None
# if model_name == "EDSR":
# model = md.load_edsr(device=device)
# print('After loading model: ')
# subprocess.run("gpustat", shell = True)
# print()
# elif model_name == "RRDB":
# model = md.load_rrdb(device=device)
# else:
# raise Exception("Unknown model...")
# model.eval()
# input_image = ut.random_image(patch_size)
# if model_name == "RRDB":
# input_image = input_image[:, 2:, :, :]
# dummy_input = input_image.to(device)
# =============================================================================
# =============================================================================
# print('Before loading model: ')
# subprocess.run("gpustat", shell = True)
# print()
# =============================================================================
device = ut.get_device_type()
model = None
if model_name == "RRDB":
model = md.load_rrdb(device)
image = ut.create_custom_npz(patch_size, patch_size)
image = image[np.newaxis, :, :]
image = torch.from_numpy(image)
dummy_input = image.unsqueeze(0).to(device)
elif model_name == "EDSR":
model = md.load_edsr(device)
# =============================================================================
# print('After loading model: ')
# subprocess.run("gpustat", shell = True)
# print()
# =============================================================================
dummy_input = ut.random_image(patch_size).to(device)
# =============================================================================
# print('After loading image: ')
# subprocess.run("gpustat", shell = True)
# print()
# =============================================================================
# =============================================================================
# b, c, h, w = 1, 3, patch_size, patch_size
# dummy_input = torch.rand(b, c, h, w, requires_grad=False).to(device)
# =============================================================================
else:
print("Unknown model!")
return
print(dummy_input.shape)
model.eval()
with torch.no_grad():
# =============================================================================
# print('Before processing: ')
# subprocess.run("gpustat", shell = True)
# print()
# =============================================================================
# =============================================================================
# output = model(input_image )
# print('After processing: ')
# subprocess.run("gpustat", shell = True)
# print()
# =============================================================================
try:
torch.onnx.export(
model,
dummy_input,
"inference_models/" + onnx_model_name,
verbose=False,
opset_version=12,
input_names=["input"],
output_names=["output"],
)
except RuntimeError as err:
sys.exit(1)
import torch
import torch.onnx
import click
import modelloader as md
import utilities as ut
import torch.nn as nn
model = md.load_edsr(device="cuda")
with torch.no_grad():
model.eval()
dummy_input = ut.random_image(100).cuda()
b, c, h, w = 1, 3, 120, 120
inp = torch.rand(b, c, h, w, requires_grad=True).cuda()
output = model(inp)
# =============================================================================
# torch.onnx.export(model,
# inp,
# "edsr.onnx",
# verbose=False,
# opset_version=11,
# input_names=['input'],
# output_names=['output'],
# dynamic_axes = {'input':{0: 'batch', 2:'height', 3:'width'},
# 'output':{0: 'batch', 2:'height', 3:'width'}
# }
# )
# =============================================================================
torch.onnx.export(
model,
inp,
"edsr_fixed.onnx",
def upsample(model_name, img_path, dimension, shave, batch_size, scale,
device):
file_name = img_path.split("/")[-1].split(".")[0]
if model_name == "RRDB":
input_image = ut.npz_loader(img_path)
c, h, w = input_image.shape
patch_list = {}
create_patch_list(patch_list, input_image, dimension, shave, scale, c,
h, w)
model = md.load_rrdb(device=device)
model.eval()
min_dim = min(dimension, h, w)
if min_dim != dimension:
print(
"\nPatch dimension is greater than the input image's minimum dimension. Changing patch dimension to input image's minimum dimension... \n "
)
dimension = min_dim
output, _ = batch_forward_chop(
patch_list,
batch_size,
c,
h,
w,
dimension,
shave,
scale,
model=model,
device=device,
)
output = output.int()
output_folder = "output_images"
ut.save_image(output,
output_folder,
h,
w,
scale,
output_file_name=file_name + "_x4")
elif model_name == "EDSR":
input_image = ut.load_image(img_path)
c, h, w = input_image.shape
patch_list = {}
create_patch_list(patch_list, input_image, dimension, shave, scale, c,
h, w)
model = md.load_edsr(device=device)
model.eval()
min_dim = min(dimension, h, w)
if min_dim != dimension:
print(
"\nPatch dimension is greater than the input image's minimum dimension. Changing patch dimension to input image's minimum dimension... \n "
)
dimension = min_dim
output, _ = batch_forward_chop(
patch_list,
batch_size,
c,
h,
w,
dimension,
shave,
scale,
model=model,
device=device,
)
# =============================================================================
# print(output)
# =============================================================================
output = output.int()
output_folder = "output_images"
ut.save_image(output,
output_folder,
h,
w,
scale,
output_file_name=file_name + "_x4")
else:
print("Unknown model...")
def patch_batch_forward_chop(
input_image,
patch_dimension,
patch_shave,
scale,
batch_size,
model_type="EDSR",
device="cuda",
print_timer=True,
):
"""
Parameters
----------
input_image : 3D Matrix
input image.
patch_dimension : int
patch dimension.
patch_shave : int
patch shave value.
scale : int
scale for LR to SR.
batch_size : int
batch size.
model_type : str, optional
model name. The default is 'EDSR'.
device : str, optional
GPU or CPU. The default is 'cuda'.
print_timer : bool, optional
print result or not. The default is True.
Returns
-------
output_image : 3D Matrix
output SR image.
"""
model = None
if model_type == "EDSR":
model = md.load_edsr(device=device)
model.eval()
elif model_type == "RRDB":
model = md.load_rrdb(device=device)
model.eval()
else:
raise Exception("{} : Unknown model...".format(model_type))
total_timer = ut.timer()
channel, height, width = input_image.shape
patch_list_timer = ut.timer()
patch_list = {}
create_patch_list(
patch_list,
input_image,
patch_dimension,
patch_shave,
scale,
channel,
height,
width,
)
patch_list_processing_time = patch_list_timer.toc()
total_batch_processing_timer = ut.timer()
output_image, timer_results = batch_forward_chop(
patch_list,
batch_size,
channel,
height,
width,
patch_dimension,
patch_shave,
scale,
model=model,
device="cuda",
print_timer=False,
)
total_batch_processing_time = total_batch_processing_timer.toc()
if model_type == "EDSR":
output_image = output_image.int()
total_time = total_timer.toc()
print(len(patch_list))
if print_timer:
print(patch_list_processing_time)
for t in timer_results:
print(t)
print(total_batch_processing_time)
print(total_time)
model = model.cpu()
del model
return output_image
def do_linear_search(test=False, test_dim=32):
"""
Linear search function...
Returns
-------
None.
"""
logger = ut.get_logger()
device = "cuda"
model_name = "EDSR"
config = toml.load("../config.toml")
run = config["run"]
scale = int(config["scale"]) if config["scale"] else 4
# device information
_, device_name = ut.get_device_details()
total, _, _ = ut.get_gpu_details(
device, "\nDevice info:", logger, print_details=False
)
log_message = (
"\nDevice: "
+ device
+ "\tDevice name: "
+ device_name
+ "\tTotal memory: "
+ str(total)
)
logger.info(log_message)
ut.clear_cuda(None, None)
state = "Before loading model: "
total, used, _ = ut.get_gpu_details(device, state, logger, print_details=True)
model = md.load_edsr(device=device)
state = "After loading model: "
total, used, _ = ut.get_gpu_details(device, state, logger, print_details=True)
# =============================================================================
# file = open("temp_max_dim.txt", "r")
# line = file.read()
# max_dim = int(line.split(":")[1])
# =============================================================================
config = toml.load("../config.toml")
max_dim = int(config["max_dim"])
if test == False:
detailed_result, memory_used, memory_free = result_from_dimension_range(
device, logger, config, model, 1, max_dim
)
else:
detailed_result, memory_used, memory_free = result_from_dimension_range(
device, logger, config, model, test_dim, test_dim
)
if test == False:
# get mean
# get std
mean_time, std_time = ut.get_mean_std(detailed_result)
mean_memory_used, std_memory_used = ut.get_mean_std(memory_used)
mean_memory_free, std_memory_free = ut.get_mean_std(memory_free)
# make folder for saving results
plt_title = "Model: {} | GPU: {} | Memory: {} MB".format(
model_name, device_name, total
)
date = "_".join(str(time.ctime()).split())
date = "_".join(date.split(":"))
foldername = date
os.mkdir("results/" + foldername)
# plot data
ut.plot_data(
foldername,
"dimension_vs_meantime",
mean_time,
"Dimensionn of Patch(nxn)",
"Mean Processing Time: LR -> SR, Scale: {} ( {} runs )".format(scale, run),
mode="mean time",
title=plt_title,
)
ut.plot_data(
foldername,
"dimension_vs_stdtime",
std_time,
"Dimension n of Patch(nxn)",
"Std of Processing Time: LR -> SR, Scale: {} ( {} runs )".format(
scale, run
),
mode="std time",
title=plt_title,
)
ut.plot_data(
foldername,
"dimension_vs_meanmemoryused",
mean_memory_used,
"Dimension n of Patch(nxn)",
"Mean Memory used: LR -> SR, Scale: {} ( {} runs )".format(scale, run),
mode="mean memory used",
title=plt_title,
)
ut.plot_data(
foldername,
"dimension_vs_stdmemoryused",
std_memory_used,
"Dimension n of Patch(nxn)",
"Std Memory Used: LR -> SR, Scale: {} ( {} runs )".format(scale, run),
mode="std memory used",
title=plt_title,
)
ut.plot_data(
foldername,
"dimension_vs_meanmemoryfree",
mean_memory_free,
"Dimension n of Patch(nxn)",
"Mean Memory Free: LR -> SR, Scale: {} ( {} runs )".format(scale, run),
mode="mean memory free",
title=plt_title,
)
ut.plot_data(
foldername,
"dimension_vs_stdmemoryfree",
std_memory_free,
"Dimension n of Patch(nxn)",
"Std Memory Free: LR -> SR, Scale: {} ( {} runs )".format(scale, run),
mode="std memory free",
title=plt_title,
)
# save data
ut.save_csv(
foldername,
"total_stat",
device,
device_name,
total,
mean_time,
std_time,
mean_memory_used,
std_memory_used,
mean_memory_free,
std_memory_free,
)