File size: 3,120 Bytes
4730cdc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Power by Zongsheng Yue 2021-11-24 20:29:36
import math
import torch
from pathlib import Path
from copy import deepcopy
from collections import OrderedDict
import torch.nn.functional as F
def calculate_parameters(net):
out = 0
for param in net.parameters():
out += param.numel()
return out
def pad_input(x, mod):
h, w = x.shape[-2:]
bottom = int(math.ceil(h/mod)*mod -h)
right = int(math.ceil(w/mod)*mod - w)
x_pad = F.pad(x, pad=(0, right, 0, bottom), mode='reflect')
return x_pad
def forward_chop(net, x, net_kwargs=None, scale=1, shave=10, min_size=160000):
n_GPUs = 1
b, c, h, w = x.size()
h_half, w_half = h // 2, w // 2
h_size, w_size = h_half + shave, w_half + shave
lr_list = [
x[:, :, 0:h_size, 0:w_size],
x[:, :, 0:h_size, (w - w_size):w],
x[:, :, (h - h_size):h, 0:w_size],
x[:, :, (h - h_size):h, (w - w_size):w]]
if w_size * h_size < min_size:
sr_list = []
for i in range(0, 4, n_GPUs):
lr_batch = torch.cat(lr_list[i:(i + n_GPUs)], dim=0)
if net_kwargs is None:
sr_batch = net(lr_batch)
else:
sr_batch = net(lr_batch, **net_kwargs)
sr_list.extend(sr_batch.chunk(n_GPUs, dim=0))
else:
sr_list = [
forward_chop(patch, shave=shave, min_size=min_size) \
for patch in lr_list
]
h, w = scale * h, scale * w
h_half, w_half = scale * h_half, scale * w_half
h_size, w_size = scale * h_size, scale * w_size
shave *= scale
output = x.new(b, c, h, w)
output[:, :, 0:h_half, 0:w_half] \
= sr_list[0][:, :, 0:h_half, 0:w_half]
output[:, :, 0:h_half, w_half:w] \
= sr_list[1][:, :, 0:h_half, (w_size - w + w_half):w_size]
output[:, :, h_half:h, 0:w_half] \
= sr_list[2][:, :, (h_size - h + h_half):h_size, 0:w_half]
output[:, :, h_half:h, w_half:w] \
= sr_list[3][:, :, (h_size - h + h_half):h_size, (w_size - w + w_half):w_size]
return output
def measure_time(net, inputs, num_forward=100):
'''
Measuring the average runing time (seconds) for pytorch.
out = net(*inputs)
'''
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
with torch.set_grad_enabled(False):
for _ in range(num_forward):
out = net(*inputs)
end.record()
torch.cuda.synchronize()
return start.elapsed_time(end) / 1000
def reload_model(model, ckpt):
module_flag = list(ckpt.keys())[0].startswith('module.')
compile_flag = '_orig_mod' in list(ckpt.keys())[0]
for source_key, source_value in model.state_dict().items():
target_key = source_key
if compile_flag and (not '_orig_mod.' in source_key):
target_key = '_orig_mod.' + target_key
if module_flag and (not source_key.startswith('module')):
target_key = 'module.' + target_key
assert target_key in ckpt
source_value.copy_(ckpt[target_key])
|