def getOptS(modelPath, modules, ramCoef): opt = Option(modelPath) weights = getStateDict(modelPath) opt.modules = modules opt.ramOffset = config.getRunType() * len(modules) for i, key in enumerate(modules): m = modules[key] wKey = m['weight'] constructor = m.get('f', 0) rc = m['ramCoef'][config.getRunType()] if 'ramCoef' in m else ramCoef[opt.ramOffset + i] o = dict((k, m[k]) for k in ('align', 'padding', 'scale') if k in m) model = initModel(opt, weights[wKey], key, constructor) if 'outShape' in m: opt.__dict__[key] = newOpt(model, rc, **o) else: model.ramCoef = rc opt.__dict__[key] = model return opt
def getOpt(optDe): model = optDe.get('model', 'dehaze') opt = Option() modelPath, opt.modelDef, ramCoef, opt.padding, opt.align, opt.prepare = mode_switch[ model] opt.ramCoef = ramCoef[config.getRunType()] opt.model = modelPath opt.modelCached = initModel(opt, modelPath, model) return opt
def newOpt(f, ramType): opt = Option() opt.modelCached = lambda x: (f(x),) opt.ramCoef = ramCoef[config.getRunType() * 2 + ramType] opt.align = 32 opt.padding = 45 opt.squeeze = identity opt.unsqueeze = identity return opt
def getOpt(model): def opt(): pass if not model in mode_switch: return opt.model = mode_switch[model][0] opt.modelDef = mode_switch[model][1] opt.ramCoef = mode_switch[model][2][config.getRunType()] opt.cropsize = config.getConfig()[1 if model[:4] == 'lite' else 2] opt.modelCached = initModel(opt, opt.model, 'DN' + model) return opt
def getOpt(model): def opt(): pass if not model in mode_switch: return opt.model = mode_switch[model][0] opt.modelDef = mode_switch[model][1] opt.ramCoef = mode_switch[model][2][config.getRunType()] opt.cropsize = config.getConfig()[1 if model[:4] == 'lite' else 2] opt.modelCached = initModel(getModel(opt), load(opt.model, map_location='cpu')) return opt
def getOpt(optDN): model = optDN['model'] if not model in mode_switch: return opt = Option(mode_switch[model][0]) opt.modelDef = mode_switch[model][1] opt.ramCoef = mode_switch[model][2][config.getRunType()] opt.cropsize = config.getConfig()[1 if model[:4] == 'lite' else 2] opt.modelCached = initModel(opt, opt.model, 'DN' + model) opt.squeeze = lambda x: x.squeeze(1) opt.unsqueeze = lambda x: x.unsqueeze(1) opt.padding = 5 return opt
def getOpt(model): def opt(): pass if not model in model_dict: return opt.model = model_dict[model] modelType = 0 if model[:4] == 'lite' else 1 opt.ramCoef = ramCoef[config.getRunType() * 2 + modelType] opt.cropsize = config.getConfig()[modelType + 1] if opt.cropsize: print('当前denoise切块大小:', opt.cropsize) opt.modelCached = getModel(opt, False) return opt
def getOpt(scale, mode, ensemble): def opt(): pass nmode = mode + str(scale) if not nmode in mode_switch: return opt.C2B = mode[:3] != 'gan' opt.mode = mode opt.model = mode_switch[nmode][0] opt.modelDef = mode_switch[nmode][1] opt.scale = scale opt.ensemble = ensemble opt.ramCoef = mode_switch[nmode][2][config.getRunType()] opt.cropsize = config.getConfig()[0] opt.modelCached = initModel(opt, opt.model, 'SR' + nmode) return opt
def getOpt(scale, mode, ensemble): def opt():pass nmode = mode+str(scale) if not nmode in mode_switch: return opt.C2B = mode[:3] != 'gan' opt.mode = mode opt.model = mode_switch[nmode][0] opt.modelDef = mode_switch[nmode][1] opt.scale = scale opt.ensemble = ensemble opt.ramCoef = mode_switch[nmode][2][config.getRunType()] opt.cropsize = config.getConfig()[0] if opt.cropsize: print('当前SR切块大小:', opt.cropsize) opt.modelCached = initModel(getModel(opt), load(opt.model, map_location='cpu')) return opt
def getOpt(optSR): opt = Option() opt.mode = optSR['model'] opt.scale = optSR['scale'] nmode = opt.mode + str(opt.scale) if not nmode in mode_switch: return if opt.mode[:3] != 'gan': opt.squeeze = lambda x: x.squeeze(1) opt.unsqueeze = lambda x: x.unsqueeze(1) opt.padding = 9 if opt.scale == 3 else 5 opt.model = mode_switch[nmode][0] opt.modelDef = mode_switch[nmode][1] opt.ensemble = optSR['ensemble'] if 'ensemble' in optSR and ( 0 <= optSR['ensemble'] <= 7) else config.ensembleSR opt.ramCoef = mode_switch[nmode][2][config.getRunType()] opt.cropsize = config.getConfig()[0] opt.modelCached = initModel(opt, opt.model, 'SR' + nmode) return opt
def getOpt(scale, mode, ensemble): def opt(): pass nmode = mode + str(scale) if not nmode in mode_switch: return opt.C2B = mode[:3] != 'gan' opt.mode = mode opt.model = mode_switch[nmode] opt.scale = scale opt.ensemble = ensemble modelType = (scale - 2) * 3 if mode[:3] == 'gan': modelType += 9 opt.ramCoef = ramCoef[config.getRunType() + modelType] opt.cropsize = config.getConfig()[0] if opt.cropsize: print('当前SR切块大小:', opt.cropsize) opt.modelCached = getModel(opt, False) return opt
def getBatchSize(option): return max(1, int((config.calcFreeMem() / option['load']) * ramCoef[config.getRunType()]))