提交 a52efec3 编写于 作者: littletomatodonkey's avatar littletomatodonkey

fix train.py

上级 093818a9
...@@ -23,7 +23,7 @@ logging.basicConfig( ...@@ -23,7 +23,7 @@ logging.basicConfig(
def time_zone(sec, fmt): def time_zone(sec, fmt):
real_time = datetime.datetime.now() + datetime.timedelta(hours=8) real_time = datetime.datetime.now()
return real_time.timetuple() return real_time.timetuple()
......
...@@ -13,12 +13,13 @@ ...@@ -13,12 +13,13 @@
# limitations under the License. # limitations under the License.
from __future__ import absolute_import from __future__ import absolute_import
import program from paddle.distributed import ParallelEnv
from ppcls.utils import logger import paddle
from ppcls.utils.save_load import init_model, save_model
from ppcls.utils.config import get_config
from ppcls.data import Reader from ppcls.data import Reader
import paddle.fluid as fluid from ppcls.utils.config import get_config
from ppcls.utils.save_load import init_model, save_model
from ppcls.utils import logger
import program
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
...@@ -53,23 +54,24 @@ def main(args): ...@@ -53,23 +54,24 @@ def main(args):
# assign the place # assign the place
use_gpu = config.get("use_gpu", True) use_gpu = config.get("use_gpu", True)
if use_gpu: if use_gpu:
gpu_id = fluid.dygraph.ParallelEnv().dev_id gpu_id = ParallelEnv().dev_id
place = fluid.CUDAPlace(gpu_id) place = paddle.CUDAPlace(gpu_id)
else: else:
place = fluid.CPUPlace() place = paddle.CPUPlace()
use_data_parallel = int(os.getenv("PADDLE_TRAINERS_NUM", 1)) != 1 use_data_parallel = int(os.getenv("PADDLE_TRAINERS_NUM", 1)) != 1
config["use_data_parallel"] = use_data_parallel config["use_data_parallel"] = use_data_parallel
with fluid.dygraph.guard(place): paddle.disable_static(place)
net = program.create_model(config.ARCHITECTURE, config.classes_num) net = program.create_model(config.ARCHITECTURE, config.classes_num)
optimizer = program.create_optimizer( optimizer = program.create_optimizer(
config, parameter_list=net.parameters()) config, parameter_list=net.parameters())
if config["use_data_parallel"]: if config["use_data_parallel"]:
strategy = fluid.dygraph.parallel.prepare_context() strategy = paddle.distributed.init_parallel_env()
net = fluid.dygraph.parallel.DataParallel(net, strategy) net = paddle.DataParallel(net, strategy)
# load model from checkpoint or pretrained model # load model from checkpoint or pretrained model
init_model(config, net, optimizer) init_model(config, net, optimizer)
...@@ -90,8 +92,7 @@ def main(args): ...@@ -90,8 +92,7 @@ def main(args):
program.run(train_dataloader, config, net, optimizer, epoch_id, program.run(train_dataloader, config, net, optimizer, epoch_id,
'train') 'train')
if not config["use_data_parallel"] or fluid.dygraph.parallel.Env( if not config["use_data_parallel"] or ParallelEnv().local_rank == 0:
).local_rank == 0:
# 2. validate with validate dataset # 2. validate with validate dataset
if config.validate and epoch_id % config.valid_interval == 0: if config.validate and epoch_id % config.valid_interval == 0:
net.eval() net.eval()
...@@ -101,15 +102,12 @@ def main(args): ...@@ -101,15 +102,12 @@ def main(args):
best_top1_acc = top1_acc best_top1_acc = top1_acc
message = "The best top1 acc {:.5f}, in epoch: {:d}".format( message = "The best top1 acc {:.5f}, in epoch: {:d}".format(
best_top1_acc, epoch_id) best_top1_acc, epoch_id)
logger.info("{:s}".format( logger.info("{:s}".format(logger.coloring(message, "RED")))
logger.coloring(message, "RED")))
if epoch_id % config.save_interval == 0: if epoch_id % config.save_interval == 0:
model_path = os.path.join( model_path = os.path.join(config.model_save_dir,
config.model_save_dir,
config.ARCHITECTURE["name"]) config.ARCHITECTURE["name"])
save_model(net, optimizer, model_path, save_model(net, optimizer, model_path, "best_model")
"best_model")
# 3. save the persistable model # 3. save the persistable model
if epoch_id % config.save_interval == 0: if epoch_id % config.save_interval == 0:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册