def defineModel(dataTup, dictP):
[nx, nxc, ny, nyc, nt, ns] = trainBasin.getSize(dataTup)
if dictP['crit'] == 'SigmaLoss':
ny = ny*2
nyc = nyc*2
# define model
if dictP['modelName'] == 'CudnnLSTM':
model = rnn.CudnnLstmModel(
nx=nx+nxc, ny=ny+nyc, hiddenSize=dictP['hiddenSize'])
elif dictP['modelName'] == 'LstmModel':
model = rnn.LstmModel(
nx=nx+nxc, ny=ny+nyc, hiddenSize=dictP['hiddenSize'])
else:
raise RuntimeError('Model not specified')
return model
def trainModelTS(outName):
outFolder = nameFolder(outName)
dictP = loadMaster(outName)
# load data
rmFlag = dictP['rmFlag'] if 'rmFlag' in dictP else False
wqData = waterQuality.DataModelWQ(dictP['dataName'], rmFlag)
varTup = (dictP['varX'], dictP['varXC'], dictP['varY'], dictP['varYC'])
dataTup, statTup = wqData.transIn(subset=dictP['trainName'], varTup=varTup)
dataTup = trainTS.dealNaN(dataTup, dictP['optNaN'])
wrapStat(outName, statTup)
# train model
[nx, nxc, ny, nyc, nt, ns] = trainTS.getSize(dataTup)
# define loss
if dictP['crit'] == 'RmseLoss':
lossFun = crit.RmseLoss()
elif dictP['crit'] == 'RmseLoss2D':
lossFun = crit.RmseLoss2D()
elif dictP['crit'] == 'SigmaLoss':
lossFun = crit.SigmaLoss()
ny = ny * 2
nyc = nyc * 2
else:
raise RuntimeError('loss function not specified')
# define model
if dictP['modelName'] == 'CudnnLSTM':
model = rnn.CudnnLstmModel(nx=nx + nxc,
ny=ny + nyc,
hiddenSize=dictP['hiddenSize'])
elif dictP['modelName'] == 'LstmModel':
model = rnn.LstmModel(nx=nx + nxc,
ny=ny + nyc,
hiddenSize=dictP['hiddenSize'])
elif dictP['modelName'] == 'AgeLSTM':
model = rnn.AgeLSTM2(nx=nx + nxc,
ny=ny,
nyc=nyc,
rho=365,
nh=dictP['hiddenSize'])
else:
raise RuntimeError('Model not specified')
if torch.cuda.is_available():
lossFun = lossFun.cuda()
model = model.cuda()
if dictP['optim'] == 'AdaDelta':
optim = torch.optim.Adadelta(model.parameters())
else:
raise RuntimeError('optimizor function not specified')
lossLst = list()
nEp = dictP['nEpoch']
sEp = dictP['saveEpoch']
logFile = os.path.join(outFolder, 'log')
if os.path.exists(logFile):
os.remove(logFile)
for k in range(0, nEp, sEp):
model, optim, lossEp = trainTS.trainModel(dataTup,
model,
lossFun,
optim,
batchSize=dictP['batchSize'],
nEp=sEp,
cEp=k,
logFile=logFile)
# save model
saveModel(outName, k + sEp, model, optim=optim)
lossLst = lossLst + lossEp
lossFile = os.path.join(outFolder, 'loss.csv')
pd.DataFrame(lossLst).to_csv(lossFile, index=False, header=False)
tempX = temp.copy()
tempX = tempX.join(dfXN)
xLst.append(tempX.values)
x = np.stack(xLst, axis=-1).swapaxes(1, 2).astype(np.float32)
y = np.stack(yLst, axis=-1).swapaxes(1, 2).astype(np.float32)
x[np.where(np.isnan(x))] = -1
ind1 = dfC.index.year <= 2016
xx = x[:, ind1, :]
yy = y[:, ind1, :]
# training
nbatch = 20
nEp = 100
ns = xx.shape[1]
nx = xx.shape[-1]
ny = yy.shape[-1]
model = rnn.LstmModel(nx=nx, ny=ny, hiddenSize=256).cuda()
optim = torch.optim.Adadelta(model.parameters())
lossFun = crit.RmseLoss().cuda()
nIterEp = int(np.ceil(np.log(0.01) / np.log(1 - nbatch / ns)))
lossLst = list()
for iEp in range(1, nEp + 1):
lossEp = 0
t0 = time.time()
for iIter in range(nIterEp):
iR = np.random.randint(0, ns, nbatch)
xTemp = xx[:, iR, :]
yTemp = yy[:, iR, :]
xT = torch.from_numpy(xTemp).float().cuda()
yT = torch.from_numpy(yTemp).float().cuda()
if iEp == 1 and iIter == 0:
try:
outName = 'weathering-FPR2QC-t365-B10'
ep = 100
# save
outFolder = basinFull.nameFolder(outName)
modelFile = os.path.join(outFolder, 'model_ep{}'.format(ep))
model = torch.load(modelFile)
modelStateFile = os.path.join(outFolder, 'modelState_ep{}'.format(ep))
torch.save(model.state_dict(), modelStateFile)
# load
dictP = basinFull.loadMaster(outName)
DF = dbBasin.DataFrameBasin(dictP['dataName'])
dictVar = {k: dictP[k] for k in ('varX', 'varXC', 'varY', 'varYC')}
DM = dbBasin.DataModelBasin(DF, subset='A10', **dictVar)
DM.loadStat(outFolder)
dataTup = DM.getData()
[nx, nxc, ny, nyc, nt, ns] = trainBasin.getSize(dataTup)
dataTup = trainBasin.dealNaN(dataTup, dictP['optNaN'])
if dictP['modelName'] == 'CudnnLSTM':
model = rnn.CudnnLstmModel(nx=nx + nxc,
ny=ny + nyc,
hiddenSize=dictP['hiddenSize'])
elif dictP['modelName'] == 'LstmModel':
model = rnn.LstmModel(nx=nx + nxc,
ny=ny + nyc,
hiddenSize=dictP['hiddenSize'])
else:
raise RuntimeError('Model not specified')
model.load_state_dict(torch.load(modelStateFile))
def trainModel(outName):
outFolder = nameFolder(outName)
dictP = loadMaster(outName)
# load data
DM = dbBasin.DataModelFull(dictP['dataName'])
varTup = (dictP['varX'], dictP['varXC'], dictP['varY'], dictP['varYC'])
dataTup = DM.extractData(varTup, dictP['subset'], dictP['sd'], dictP['ed'])
if dictP['borrowStat'] is None:
dataTup, statTup = DM.transIn(dataTup, varTup)
else:
statTup = loadStat(dictP['borrowStat'])
dataTup = DM.transIn(dataTup, varTup, statTup=statTup)
dataTup = trainBasin.dealNaN(dataTup, dictP['optNaN'])
wrapStat(outName, statTup)
# train model
[nx, nxc, ny, nyc, nt, ns] = trainBasin.getSize(dataTup)
# define loss
lossFun = getattr(crit, dictP['crit'])()
if dictP['crit'] == 'SigmaLoss':
ny = ny * 2
nyc = nyc * 2
# define model
if dictP['modelName'] == 'CudnnLSTM':
model = rnn.CudnnLstmModel(nx=nx + nxc,
ny=ny + nyc,
hiddenSize=dictP['hiddenSize'])
elif dictP['modelName'] == 'LstmModel':
model = rnn.LstmModel(nx=nx + nxc,
ny=ny + nyc,
hiddenSize=dictP['hiddenSize'])
else:
raise RuntimeError('Model not specified')
if torch.cuda.is_available():
lossFun = lossFun.cuda()
model = model.cuda()
if dictP['optim'] == 'AdaDelta':
optim = torch.optim.Adadelta(model.parameters())
else:
raise RuntimeError('optimizor function not specified')
lossLst = list()
nEp = dictP['nEpoch']
sEp = dictP['saveEpoch']
logFile = os.path.join(outFolder, 'log')
if os.path.exists(logFile):
os.remove(logFile)
for k in range(0, nEp, sEp):
model, optim, lossEp = trainBasin.trainModel(
dataTup,
model,
lossFun,
optim,
batchSize=dictP['batchSize'],
nEp=sEp,
cEp=k,
logFile=logFile)
# save model
saveModel(outName, k + sEp, model, optim=optim)
lossLst = lossLst + lossEp
lossFile = os.path.join(outFolder, 'loss.csv')
pd.DataFrame(lossLst).to_csv(lossFile, index=False, header=False)