def initModel(model, weights=None):
if weights:
print('reloading weights')
model.load_state_dict(weights)
for param in model.parameters():
param.requires_grad_(False)
return model.eval().to(dtype=config.dtype(), device=config.device())
def initModel(opt, weights=None, key=None, f=lambda opt: opt.modelDef()):
if key and key in modelCache:
return modelCache[key].to(dtype=config.dtype(), device=config.device())
log.info('loading model {}'.format(opt.model))
model = f(opt)
if weights:
log.info('reloading weights')
if type(weights) == str:
weights = getStateDict(weights)
model.load_state_dict(weights)
for param in model.parameters():
param.requires_grad_(False)
model.eval()
if key:
modelCache[key] = model
return model.to(dtype=config.dtype(), device=config.device())
def procInput(source, bitDepth, fs, out):
out['load'] = 1
node = Node({'op': 'toTorch', 'bits': bitDepth})
fs.append(
NonNullWrap(
node.bindFunc(toTorch(bitDepth, config.dtype(), config.device()))))
return fs, [node], out
def prepare():
device = config.device()
model = getModel(device)
extract = getExtract(device)
data = DataGenerator(config.imageDir)
targetStyleTensor = imageTensor(config.styleImage).to(device)
targetStyleFeature = extract(targetStyleTensor)
styleLossFunc = createStyleLossFunc(targetStyleFeature, config.netStyleWeight)
optimizer = config.optimizer(model)
return device, model, optimizer, extract, styleLossFunc, data
def getModel(opt, cache=True):
if hasattr(opt, 'modelCached') and cache:
return opt.modelCached
print('loading model {}'.format(opt.model))
model = f(opt)
print('reloading weights')
weights = torch.load(opt.model)
model.load_state_dict(weights)
model.eval().to(dtype=config.dtype(), device=config.device())
for param in model.parameters():
param.requires_grad_(False)
return model
def prepare(shape, ram, ramCoef, pad, sc, align=8, cropsize=0):
*_, c, h, w = shape
n = solveRam(ram, c, ramCoef)
af = alignF[align]
s = af(minSize + pad * 2)
if n < s * s:
memoryError(ram)
ph, pw = max(1, h - pad * 3), max(1, w - pad * 3)
ns = np.arange(s / align, int(n / (align * s)) + 1, dtype=np.int)
ms = (n / (align * align) / ns).astype(int)
ns, ms = ns * align, ms * align
nn, mn = np.ceil(ph / (ns - 2 * pad)).clip(2), np.ceil(
pw / (ms - 2 * pad)).clip(2)
nn[ns >= h] = 1
mn[ms >= w] = 1
ds = nn * mn # minimize number of clips
ind = np.argwhere(ds == ds.min()).squeeze(1)
mina = ind[np.abs(ind - len(ds) / 2).argmin(
)] # pick the size with ratio of width and height closer to 1
ah, aw, acs = af(h), af(w), af(cropsize)
ih, iw = (min(acs, ns[mina]),
min(acs, ms[mina])) if cropsize > 0 else (ns[mina], ms[mina])
ih, iw = min(ah, ih), min(aw, iw)
startH, endH, clipH, stepH, bH = getAnchors(h, ph, ih, pad, af, sc)
startW, endW, clipW, stepW, wH = getAnchors(w, pw, iw, pad, af, sc)
padSc, outh, outw = pad * sc, h * sc, w * sc
if (stepH > 1) and (stepW > 1):
padImage = identity
unpad = identity
elif stepH > 1:
padImage = getPad(aw, w, 0, 0)
unpad = lambda im: im[:, :, :outw]
elif stepW > 1:
padImage = getPad(0, 0, ah, h)
unpad = lambda im: im[:, :outh]
else:
padImage = getPad(aw, w, ah, h)
unpad = lambda im: im[:, :outh, :outw]
b = ((torch.arange(padSc, dtype=config.dtype(), device=config.device()) /
padSc - .5) * 9).sigmoid().view(1, -1)
def iterClip():
for i in range(stepH):
top, bottom, bsc = startH[i], endH[i], bH[i]
topT = clipH if i == stepH - 1 else (0 if i == 0 else padSc)
for j in range(stepW):
left, right, rsc = startW[j], endW[j], wH[j]
leftT = clipW if j == stepW - 1 else (0 if j == 0 else padSc)
yield (top, bottom, left, right, topT, leftT, bsc, rsc)
return iterClip, padImage, unpad, (*shape[:-2], outh, outw), b
def __init__(self, window=None, device=config.device(), offload=True, store=True, tensor=True, name=None, reserve=0, **_): self.source = DefaultStreamSource() next(self.source) self.wm1 = window - 1 if window else 0 self.device = device self.tensor = tensor self.offload = offload and store self.store = store self.batchFunc = torch.stack if tensor else identity self.name = name self.start = 0 self.end = 0 self.state = [] self.reserve = reserve # ensure enough items to pad self.stateR = []
def train(contentImage, styleImage, reCreate=None):
device = config.device()
contentImage = contentImage.to(device)
styleImage = styleImage.to(device)
imshow = None
if config.show_interval > 0:
imshow = createShowImageFunc(True)
if reCreate is None:
reCreate = contentImage.clone()
extract = getExtract(device)
originContentFeature = extract(contentImage)
originStyleFeature = extract(styleImage)
# 梯度训练的是reCreate
optimizer = optim.LBFGS([reCreate.requires_grad_()])
lossFunc = createLossFunc(originContentFeature, originStyleFeature,
config.singletonStyleWeight,
config.singletonContentWeight)
iteration = [0]
newSaveFileFunc = createSaveFileNameFunc(config.contentImage,
config.create_dir)
trainLogger = createLogger(True)
def singletonTraining():
optimizer.zero_grad()
createFeature = extract(reCreate)
styleLoss, contentLoss = lossFunc(createFeature)
loss = styleLoss + contentLoss
if iteration[0] % config.save_interval == 0:
saveTensorImage(reCreate, newSaveFileFunc(iteration[0]))
if (imshow is not None) and (iteration[0] % config.show_interval == 0):
imshow(reCreate, iteration[0])
if iteration[0] % config.log_interval == 0:
trainLogger(
'iter : {:4.0f} \tcontentLoss : {:10.5f} \tstyleLoss : {:10.5f}'
.format(iteration[0], contentLoss, styleLoss))
# 单次训练需要记住上次的值，否则会释放
loss.backward(retain_graph=True)
iteration[0] += 1
return loss
# 梯度回传的就是张量本身而非网络结构
# 需要对单一的相同的张量进行重复计算
while iteration[0] < config.iterations:
optimizer.step(singletonTraining)
return reCreate
def genProcess(scale=1,
mode='a',
dnmodel='no',
dnseq='before',
source='image',
bitDepthIn=8,
bitDepthOut=0):
SRopt = runSR.getOpt(scale, mode, config.ensembleSR)
DNopt = runDN.getOpt(dnmodel)
config.getFreeMem(True)
if not bitDepthOut:
bitDepthOut = bitDepthIn
quant = 1 << bitDepthIn
funcs = []
last = lambda im, _: im
if source == 'file':
funcs.append(readFile)
elif source == 'buffer':
funcs.append(toNumPy(bitDepthIn))
funcs.append(toTorch(quant, config.dtype(), config.device()))
if (dnseq == 'before') and (dnmodel != 'no'):
funcs.append(lambda im: runDN.dn(im, DNopt))
if (scale > 1):
funcs.append(lambda im: runSR.sr(im, SRopt))
if (dnseq == 'after') and (dnmodel != 'no'):
funcs.append(lambda im: runDN.dn(im, DNopt))
funcs.append(toOutput(bitDepthOut))
if source == 'file':
last = writeFile
elif source == 'buffer':
funcs.append(toBuffer(bitDepthOut))
def process(im, name=None):
im = reduce(apply, funcs, im)
return last(im, name)
return process
import sys
sys.path.append('./python')
import torch
from procedure import genProcess
from config import config, process
getMemUsed = torch.cuda.max_memory_cached if config.cuda else lambda: process.memory_info()[0]
step = dict(op='dehaze', model='sun')
load = 3 << 20
p, _ = genProcess([step], True, dict(bitDepth=8, channel=0, source=0, load=load))
t = torch.randn((3, 1024, 1024), dtype=config.dtype(), device=config.device()) # pylint: disable=E1101
m = getMemUsed()
print(config.dtype(), m)
sys.stdin.readline()
if config.cuda:
torch.cuda.reset_max_memory_cached()
p(t)
m = getMemUsed() if config.cuda else (getMemUsed() - m)
print(m, m / load)
sys.stdin.readline()
def getFlowBack(opt, width, height):
if opt.flowBackWarp:
return opt.flowBackWarp
opt.flowBackWarp = initModel(backWarp(width, height, config.device(), config.dtype()))
return opt.flowBackWarp
sys.path.append('./python')
from time import perf_counter
import torch
from torch.profiler import profile, schedule, ProfilerActivity
from config import config
from imageProcess import doCrop
from videoSR import getOpt
modelConfig = None
shape = (1, 2, 3, 1088, 1920)
p = getOpt(modelConfig).spynet#.modelCached
#p.outShape = (1, 64, 1088, 1920)
getMemUsed = lambda i: torch.cuda.memory_stats(i)['reserved_bytes.all.peak']
t = torch.randn(shape, dtype=config.dtype(), device=config.device()) # pylint: disable=E1101
load = shape[-1] * shape[-2] * shape[0]
m = getMemUsed(config.device()) if config.cuda else None
print(config.dtype(), config.device(), m)
if config.cuda:
p(t)
#doCrop(p, t)
getMemUsed(config.device())
start = perf_counter()
p(t)
#doCrop(p, t).mean().cpu()
print('time elpased: {}'.format(perf_counter() - start))
m = getMemUsed(config.device())
else:
schedule1 = schedule(
wait=1,
# coding: utf-8
# In[1]:
import torch
import sys
sys.path.append('./python')
from progress import initialETA, ops, newOp, slideAverage
from imageProcess import genProcess
from config import config
inTensor = torch.randn((3, 1080, 1920),
dtype=config.dtype(),
device=config.device())
imgType = dict(bitDepth=8, channel=0, source=0, load=inTensor.nelement())
opt1 = dict(op='SR', model='a', scale=2)
opt2 = dict(op='SR', model='lite', scale=2)
def run(cases, times=1):
for _, node, _1 in cases:
node.learn = times
ops[node.op] = newOp(slideAverage(1 - 1 / times))
for i in range(times):
for process, node, _ in cases:
initialETA(node)
process(inTensor)
torch.cuda.empty_cache()
inputFolder = '../test-pics'
refFile = 0 #'test/1566005911.7879605_ci.png'
def context():
pass
opt = Option(
('test/{}.pth' if test else 'model/demoire/{}.pth').format(modelName))
opt.padding = 31
opt.ramCoef = 1 / 8000.
opt.align = 128
opt.modelCached = initModel(opt, weights=opt.model, f=lambda _: Net())
toTorch = lambda x: torch.from_numpy(np.array(x)).permute(2, 0, 1).to(
dtype=config.dtype(), device=config.device()) / 256
time = 0.0
for pic in os.listdir(inputFolder):
original = toTorch(readFile(context=context)(inputFolder + '/' + pic))
ref = toTorch(readFile(context=context)(refFile + '/' +
pic)) if refFile else original
start = perf_counter()
y = ensemble(opt)(original)
time += perf_counter() - start
print(pic, float(y.mean(dtype=torch.float)),
float((y - ref).abs().mean(dtype=torch.float)))
out = toOutput(8)(toFloat(y))
writeFile(
out,
'download/{}.{}.png'.format(splitext(split(pic)[1])[0],
modelName), context)
#from torchvision.transforms import Normalize
from slomo import UNet, backWarp
from imageProcess import initModel, getStateDict
from config import config
log = logging.getLogger('Moe')
modelPath = './model/slomo/SuperSloMo.ckpt'
ramCoef = [.9 / x for x in (6418.7, 1393., 1156.3)]
#mean = [0.429, 0.431, 0.397]
#std = [1, 1, 1]
#negMean = [-x for x in mean]
#identity = lambda x, *_: x
upTruncBy32 = lambda x: (-x & 0xffffffe0 ^ 0xffffffff) + 1
getFlowComp = lambda *_: UNet(6, 4)
getFlowIntrp = lambda *_: UNet(20, 5)
getFlowBack = lambda opt: backWarp(opt.width, opt.height, config.device(), config.dtype())
def getOpt(option):
def opt():pass
# Initialize model
opt.model = modelPath
dict1 = getStateDict(modelPath)
opt.flowComp = initModel(opt, dict1['state_dictFC'], 'flowComp', getFlowComp)
opt.ArbTimeFlowIntrp = initModel(opt, dict1['state_dictAT'], 'ArbTimeFlowIntrp', getFlowIntrp)
opt.sf = option['sf']
opt.firstTime = 1
opt.notLast = 1
opt.batchSize = 0
if opt.sf < 2:
raise RuntimeError('Error: --sf/slomo factor has to be at least 2')
return opt
import sys
sys.path.append('./python')
import torch
from imageProcess import genProcess
from config import config
p, _ = genProcess([dict(op='SR', scale=2, model='lite')], True, dict(bitDepth=8, channel=0, source=0, load=1 << 16))
t = torch.randn((1, 510, 510), dtype=config.dtype(), device=config.device()) # pylint: disable=E1101
p(t)
