December 2021 update

Author: liuqiuhui2015

adv/train/mulang/train_m2o.py

@@ -11,8 +11,10 @@
 from utils.base import *
 from utils.init import init_model_params
 from utils.contpara import get_model_parameters
-from utils.h5serial import h5save, h5load
-from utils.fmt.base import tostr, save_states, load_states, pad_id
+from utils.state.holder import Holder
+from utils.state.pyrand import PyRandomState
+from utils.state.thrand import THRandomState
+from utils.fmt.base import tostr, pad_id
 from utils.fmt.base4torch import parse_cuda, load_emb
 from utils.mulang import data_sampler
 
@@ -30,7 +32,7 @@
 
 from transformer.NMT import NMT
 
-def train(td, tl, ed, nd, optm, lrsch, model, lossf, mv_device, logger, done_tokens, multi_gpu, multi_gpu_optimizer, tokens_optm=32768, nreport=None, save_every=None, chkpf=None, chkpof=None, statesf=None, num_checkpoint=1, cur_checkid=0, report_eva=True, remain_steps=None, save_loss=False, save_checkp_epoch=False, scaler=None):
+def train(td, tl, ed, nd, optm, lrsch, model, lossf, mv_device, logger, done_tokens, multi_gpu, multi_gpu_optimizer, tokens_optm=32768, nreport=None, save_every=None, chkpf=None, state_holder=None, statesf=None, num_checkpoint=1, cur_checkid=0, report_eva=True, remain_steps=None, save_loss=False, save_checkp_epoch=False, scaler=None):
 
 	sum_loss = part_loss = 0.0
 	sum_wd = part_wd = 0
@@ -77,17 +79,12 @@ def train(td, tl, ed, nd, optm, lrsch, model, lossf, mv_device, logger, done_tok
 					if num_checkpoint > 1:
 						_fend = "_%d.h5" % (_cur_checkid)
 						_chkpf = chkpf[:-3] + _fend
-						if chkpof is not None:
-							_chkpof = chkpof[:-3] + _fend
 						_cur_checkid = (_cur_checkid + 1) % num_checkpoint
 					else:
 						_chkpf = chkpf
-						_chkpof = chkpof
 					save_model(model, _chkpf, multi_gpu, print_func=logger.info)
-					if chkpof is not None:
-						h5save(optm.state_dict(), _chkpof)
 					if statesf is not None:
-						save_states(statesf, tl[cur_b - 1:])
+						save_states(state_holder.state_dict(update=False, **{"remain_steps": _cur_rstep, "checkpoint_id": _cur_checkid, "training_list": tl[cur_b - 1:]}), statesf, print_func=logger.info)
 				_cur_rstep -= 1
 				if _cur_rstep <= 0:
 					break
@@ -111,17 +108,12 @@ def train(td, tl, ed, nd, optm, lrsch, model, lossf, mv_device, logger, done_tok
 			if num_checkpoint > 1:
 				_fend = "_%d.h5" % (_cur_checkid)
 				_chkpf = chkpf[:-3] + _fend
-				if chkpof is not None:
-					_chkpof = chkpof[:-3] + _fend
 				_cur_checkid = (_cur_checkid + 1) % num_checkpoint
 			else:
 				_chkpf = chkpf
-				_chkpof = chkpof
 			save_model(model, _chkpf, multi_gpu, print_func=logger.info)
-			if chkpof is not None:
-				h5save(optm.state_dict(), _chkpof)
 			if statesf is not None:
-				save_states(statesf, tl[cur_b - 1:])
+				save_states(state_holder.state_dict(update=False, **{"remain_steps": _cur_rstep, "checkpoint_id": _cur_checkid, "training_list": tl[cur_b - 1:]}), statesf, print_func=logger.info)
 		cur_b += 1
 	if part_wd != 0.0:
 		logger.info("Average loss over %d tokens: %.3f" % (part_wd, part_loss / part_wd,))
@@ -181,7 +173,7 @@ def load_fixing(module):
 batch_report = cnfg.batch_report
 report_eva = cnfg.report_eva
 use_ams = cnfg.use_ams
-save_optm_state = cnfg.save_optm_state
+cnt_states = cnfg.train_statesf
 save_auto_clean = cnfg.save_auto_clean
 overwrite_eva = cnfg.overwrite_eva
 save_every = cnfg.save_every
@@ -193,14 +185,11 @@ def load_fixing(module):
 mkdir(wkdir)
 
 chkpf = None
-chkpof = None
 statesf = None
 if save_every is not None:
 	chkpf = wkdir + "checkpoint.h5"
-	if save_optm_state:
-		chkpof = wkdir + "checkpoint.optm.h5"
-	if cnfg.save_train_state:
-		statesf = wkdir + "checkpoint.states"
+if cnfg.save_train_state:
+	statesf = wkdir + "train.states.t7"
 
 logger = get_logger(wkdir + "train.log")
 
@@ -217,10 +206,6 @@ def load_fixing(module):
 nword = td["nword"][:].tolist()
 nwordi, ntask, nwordt = nword[0], nword[1], nword[-1]
 
-logger.info("Design models with seed: %d" % torch.initial_seed())
-mymodel = NMT(cnfg.isize, nwordi, nwordt, cnfg.nlayer, cnfg.ff_hsize, cnfg.drop, cnfg.attn_drop, cnfg.share_emb, cnfg.nhead, cache_len_default, cnfg.attn_hsize, cnfg.norm_output, cnfg.bindDecoderEmb, cnfg.forbidden_indexes)
-
-fine_tune_m = cnfg.fine_tune_m
 task_weight, task_weight_T = cnfg.task_weight, cnfg.task_weight_T
 if task_weight_T is None or task_weight_T == 1.0:
 	tl = [(str(i), _task,) for _nd, _task in zip(ntrain, td["taskorder"][:].tolist()) for i in range(_nd)]
@@ -234,6 +219,11 @@ def load_fixing(module):
 	train_sampler = data_sampler(ntrain if task_weight is None else task_weight, task_weight_T, ntrain, train_taskorder, nsample=sum(ntrain))
 nvalid = [(str(i), _task,) for _nd, _task in zip(nvalid, vd["taskorder"][:].tolist()) for i in range(_nd)]
 
+logger.info("Design models with seed: %d" % torch.initial_seed())
+mymodel = NMT(cnfg.isize, nwordi, nwordt, cnfg.nlayer, cnfg.ff_hsize, cnfg.drop, cnfg.attn_drop, cnfg.share_emb, cnfg.nhead, cache_len_default, cnfg.attn_hsize, cnfg.norm_output, cnfg.bindDecoderEmb, cnfg.forbidden_indexes)
+
+fine_tune_m = cnfg.fine_tune_m
+
 mymodel = init_model_params(mymodel)
 mymodel.apply(init_fixing)
 if fine_tune_m is not None:
@@ -267,13 +257,10 @@ def load_fixing(module):
 	optimizer = Optimizer(get_model_parameters(mymodel, contiguous_parameters=contiguous_parameters), lr=init_lr, betas=adam_betas_default, eps=ieps_adam_default, weight_decay=cnfg.weight_decay, amsgrad=use_ams)
 optimizer.zero_grad(set_to_none=optm_step_zero_grad_set_none)
 
-fine_tune_state = cnfg.fine_tune_state
-if fine_tune_state is not None:
-	logger.info("Load optimizer state from: " + fine_tune_state)
-	optimizer.load_state_dict(h5load(fine_tune_state))
-
 lrsch = LRScheduler(optimizer, cnfg.isize, cnfg.warm_step, scale=cnfg.lr_scale)
 
+state_holder = None if statesf is None and cnt_states is None else Holder(**{"optm": optimizer, "lrsch": lrsch, "pyrand": PyRandomState(), "thrand": THRandomState(use_cuda=use_cuda)})
+
 num_checkpoint = cnfg.num_checkpoint
 cur_checkid = 0
 
@@ -286,15 +273,22 @@ def load_fixing(module):
 	save_model(mymodel, wkdir + "init.h5", multi_gpu, print_func=logger.info)
 	logger.info("Initial model saved")
 else:
-	cnt_states = cnfg.train_statesf
 	if cnt_states is not None:
-		logger.info("Continue last epoch")
-		tminerr, done_tokens, cur_checkid, remain_steps, _ = train(td, load_states(cnt_states), vd, nvalid, optimizer, lrsch, mymodel, lossf, cuda_device, logger, done_tokens, multi_gpu, multi_gpu_optimizer, tokens_optm, batch_report, save_every, chkpf, chkpof, statesf, num_checkpoint, cur_checkid, report_eva, remain_steps, False, False, scaler)
+		logger.info("Loading training states")
+		_remain_states = state_holder.load_state_dict(torch.load(cnt_states))
+		remain_steps, cur_checkid = _remain_states["remain_steps"], _remain_states["checkpoint_id"]
+		if "training_list" in _remain_states:
+			_ctl = _remain_states["training_list"]
+		else:
+			shuffle(tl)
+			_ctl = tl
+		tminerr, done_tokens, cur_checkid, remain_steps, _ = train(td, _ctl, vd, nvalid, optimizer, lrsch, mymodel, lossf, cuda_device, logger, done_tokens, multi_gpu, multi_gpu_optimizer, tokens_optm, batch_report, save_every, chkpf, state_holder, statesf, num_checkpoint, cur_checkid, report_eva, remain_steps, False, False, scaler)
+		_ctl = _remain_states = None
 		vloss, vprec = eva(vd, nvalid, mymodel, lossf, cuda_device, multi_gpu, use_amp)
 		logger.info("Epoch: 0, train loss: %.3f, valid loss/error: %.3f %.2f" % (tminerr, vloss, vprec,))
 		save_model(mymodel, wkdir + "train_0_%.3f_%.3f_%.2f.h5" % (tminerr, vloss, vprec,), multi_gpu, print_func=logger.info, mtyp=("eva" if overwrite_eva else "train") if save_auto_clean else None)
-		if save_optm_state:
-			h5save(optimizer.state_dict(), wkdir + "train_0_%.3f_%.3f_%.2f.optm.h5" % (tminerr, vloss, vprec,))
+		if statesf is not None:
+			save_states(state_holder.state_dict(update=False, **{"remain_steps": remain_steps, "checkpoint_id": cur_checkid}), statesf, print_func=logger.info)
 		logger.info("New best model saved")
 
 if cnfg.dss_ws is not None and cnfg.dss_ws > 0.0 and cnfg.dss_ws < 1.0:
@@ -319,14 +313,14 @@ def load_fixing(module):
 	else:
 		tl = train_sampler.generate()
 	free_cache(use_cuda)
-	terr, done_tokens, cur_checkid, remain_steps, _Dws = train(td, tl, vd, nvalid, optimizer, lrsch, mymodel, lossf, cuda_device, logger, done_tokens, multi_gpu, multi_gpu_optimizer, tokens_optm, batch_report, save_every, chkpf, chkpof, statesf, num_checkpoint, cur_checkid, report_eva, remain_steps, dss_ws > 0, i >= start_chkp_save, scaler)
+	terr, done_tokens, cur_checkid, remain_steps, _Dws = train(td, tl, vd, nvalid, optimizer, lrsch, mymodel, lossf, cuda_device, logger, done_tokens, multi_gpu, multi_gpu_optimizer, tokens_optm, batch_report, save_every, chkpf, state_holder, statesf, num_checkpoint, cur_checkid, report_eva, remain_steps, dss_ws > 0, i >= start_chkp_save, scaler)
 	vloss, vprec = eva(vd, nvalid, mymodel, lossf, cuda_device, multi_gpu, use_amp)
 	logger.info("Epoch: %d, train loss: %.3f, valid loss/error: %.3f %.2f" % (i, terr, vloss, vprec,))
 
 	if (vprec <= minerr) or (vloss <= minloss):
 		save_model(mymodel, wkdir + "eva_%d_%.3f_%.3f_%.2f.h5" % (i, terr, vloss, vprec,), multi_gpu, print_func=logger.info, mtyp="eva" if save_auto_clean else None)
-		if save_optm_state:
-			h5save(optimizer.state_dict(), wkdir + "eva_%d_%.3f_%.3f_%.2f.optm.h5" % (i, terr, vloss, vprec,))
+		if statesf is not None:
+			save_states(state_holder.state_dict(update=False, **{"remain_steps": remain_steps, "checkpoint_id": cur_checkid}), statesf, print_func=logger.info)
 		logger.info("New best model saved")
 
 		namin = 0
@@ -340,15 +334,18 @@ def load_fixing(module):
 		if terr < tminerr:
 			tminerr = terr
 			save_model(mymodel, wkdir + "train_%d_%.3f_%.3f_%.2f.h5" % (i, terr, vloss, vprec,), multi_gpu, print_func=logger.info, mtyp=("eva" if overwrite_eva else "train") if save_auto_clean else None)
-			if save_optm_state:
-				h5save(optimizer.state_dict(), wkdir + "train_%d_%.3f_%.3f_%.2f.optm.h5" % (i, terr, vloss, vprec,))
+			if statesf is not None:
+				save_states(state_holder.state_dict(update=False, **{"remain_steps": remain_steps, "checkpoint_id": cur_checkid}), statesf, print_func=logger.info)
 		elif epoch_save:
 			save_model(mymodel, wkdir + "epoch_%d_%.3f_%.3f_%.2f.h5" % (i, terr, vloss, vprec,), multi_gpu, print_func=logger.info)
+			if statesf is not None:
+				save_states(state_holder.state_dict(update=False, **{"remain_steps": remain_steps, "checkpoint_id": cur_checkid}), statesf, print_func=logger.info)
 
 		namin += 1
 		if namin >= earlystop:
 			if done_tokens > 0:
 				optm_step(optimizer, model=mymodel, scaler=scaler, multi_gpu=multi_gpu, multi_gpu_optimizer=multi_gpu_optimizer)
+				lrsch.step()
 				done_tokens = 0
 			logger.info("early stop")
 			break
@@ -368,10 +365,11 @@ def load_fixing(module):
 
 if done_tokens > 0:
 	optm_step(optimizer, model=mymodel, scaler=scaler, multi_gpu=multi_gpu, multi_gpu_optimizer=multi_gpu_optimizer)
+	lrsch.step()
 
 save_model(mymodel, wkdir + "last.h5", multi_gpu, print_func=logger.info)
-if save_optm_state:
-	h5save(optimizer.state_dict(), wkdir + "last.optm.h5")
+if statesf is not None:
+	save_states(state_holder.state_dict(update=False, **{"remain_steps": remain_steps, "checkpoint_id": cur_checkid}), statesf, print_func=logger.info)
 logger.info("model saved")
 
 td.close()