This commit is contained in:
494
test_tipc/supplementary/train.py
Normal file
494
test_tipc/supplementary/train.py
Normal file
@@ -0,0 +1,494 @@
|
||||
import paddle
|
||||
import numpy as np
|
||||
import os
|
||||
import paddle.nn as nn
|
||||
import paddle.distributed as dist
|
||||
|
||||
dist.get_world_size()
|
||||
dist.init_parallel_env()
|
||||
|
||||
from loss import build_loss, LossDistill, DMLLoss, KLJSLoss
|
||||
from optimizer import create_optimizer
|
||||
from data_loader import build_dataloader
|
||||
from metric import create_metric
|
||||
from mv3 import MobileNetV3_large_x0_5, distillmv3_large_x0_5, build_model
|
||||
from config import preprocess
|
||||
import time
|
||||
|
||||
from paddleslim.dygraph.quant import QAT
|
||||
from slim.slim_quant import PACT, quant_config
|
||||
from slim.slim_fpgm import prune_model
|
||||
from utils import load_model
|
||||
|
||||
|
||||
def _mkdir_if_not_exist(path, logger):
|
||||
"""
|
||||
mkdir if not exists, ignore the exception when multiprocess mkdir together
|
||||
"""
|
||||
if not os.path.exists(path):
|
||||
try:
|
||||
os.makedirs(path)
|
||||
except OSError as e:
|
||||
if e.errno == errno.EEXIST and os.path.isdir(path):
|
||||
logger.warning(
|
||||
"be happy if some process has already created {}".format(path)
|
||||
)
|
||||
else:
|
||||
raise OSError("Failed to mkdir {}".format(path))
|
||||
|
||||
|
||||
def save_model(
|
||||
model, optimizer, model_path, logger, is_best=False, prefix="ppocr", **kwargs
|
||||
):
|
||||
"""
|
||||
save model to the target path
|
||||
"""
|
||||
_mkdir_if_not_exist(model_path, logger)
|
||||
model_prefix = os.path.join(model_path, prefix)
|
||||
paddle.save(model.state_dict(), model_prefix + ".pdparams")
|
||||
if type(optimizer) is list:
|
||||
paddle.save(optimizer[0].state_dict(), model_prefix + ".pdopt")
|
||||
paddle.save(optimizer[1].state_dict(), model_prefix + "_1" + ".pdopt")
|
||||
|
||||
else:
|
||||
paddle.save(optimizer.state_dict(), model_prefix + ".pdopt")
|
||||
|
||||
# # save metric and config
|
||||
# with open(model_prefix + '.states', 'wb') as f:
|
||||
# pickle.dump(kwargs, f, protocol=2)
|
||||
if is_best:
|
||||
logger.info("save best model is to {}".format(model_prefix))
|
||||
else:
|
||||
logger.info("save model in {}".format(model_prefix))
|
||||
|
||||
|
||||
def amp_scaler(config):
|
||||
if "AMP" in config and config["AMP"]["use_amp"] is True:
|
||||
AMP_RELATED_FLAGS_SETTING = {
|
||||
"FLAGS_cudnn_batchnorm_spatial_persistent": 1,
|
||||
}
|
||||
paddle.set_flags(AMP_RELATED_FLAGS_SETTING)
|
||||
scale_loss = config["AMP"].get("scale_loss", 1.0)
|
||||
use_dynamic_loss_scaling = config["AMP"].get("use_dynamic_loss_scaling", False)
|
||||
scaler = paddle.amp.GradScaler(
|
||||
init_loss_scaling=scale_loss,
|
||||
use_dynamic_loss_scaling=use_dynamic_loss_scaling,
|
||||
)
|
||||
return scaler
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def set_seed(seed):
|
||||
paddle.seed(seed)
|
||||
np.random.seed(seed)
|
||||
|
||||
|
||||
def train(config, scaler=None):
|
||||
EPOCH = config["epoch"]
|
||||
topk = config["topk"]
|
||||
|
||||
batch_size = config["TRAIN"]["batch_size"]
|
||||
num_workers = config["TRAIN"]["num_workers"]
|
||||
train_loader = build_dataloader(
|
||||
"train", batch_size=batch_size, num_workers=num_workers
|
||||
)
|
||||
|
||||
# build metric
|
||||
metric_func = create_metric
|
||||
|
||||
# build model
|
||||
# model = MobileNetV3_large_x0_5(class_dim=100)
|
||||
model = build_model(config)
|
||||
|
||||
# build_optimizer
|
||||
optimizer, lr_scheduler = create_optimizer(
|
||||
config, parameter_list=model.parameters()
|
||||
)
|
||||
|
||||
# load model
|
||||
pre_best_model_dict = load_model(config, model, optimizer)
|
||||
if len(pre_best_model_dict) > 0:
|
||||
pre_str = "The metric of loaded metric as follows {}".format(
|
||||
", ".join(["{}: {}".format(k, v) for k, v in pre_best_model_dict.items()])
|
||||
)
|
||||
logger.info(pre_str)
|
||||
|
||||
# about slim prune and quant
|
||||
if "quant_train" in config and config["quant_train"] is True:
|
||||
quanter = QAT(config=quant_config, act_preprocess=PACT)
|
||||
quanter.quantize(model)
|
||||
elif "prune_train" in config and config["prune_train"] is True:
|
||||
model = prune_model(model, [1, 3, 32, 32], 0.1)
|
||||
else:
|
||||
pass
|
||||
|
||||
# distribution
|
||||
model.train()
|
||||
model = paddle.DataParallel(model)
|
||||
# build loss function
|
||||
loss_func = build_loss(config)
|
||||
|
||||
data_num = len(train_loader)
|
||||
|
||||
best_acc = {}
|
||||
for epoch in range(EPOCH):
|
||||
st = time.time()
|
||||
for idx, data in enumerate(train_loader):
|
||||
img_batch, label = data
|
||||
img_batch = paddle.transpose(img_batch, [0, 3, 1, 2])
|
||||
label = paddle.unsqueeze(label, -1)
|
||||
|
||||
if scaler is not None:
|
||||
with paddle.amp.auto_cast():
|
||||
outs = model(img_batch)
|
||||
else:
|
||||
outs = model(img_batch)
|
||||
|
||||
# cal metric
|
||||
acc = metric_func(outs, label)
|
||||
|
||||
# cal loss
|
||||
avg_loss = loss_func(outs, label)
|
||||
|
||||
if scaler is None:
|
||||
# backward
|
||||
avg_loss.backward()
|
||||
optimizer.step()
|
||||
optimizer.clear_grad()
|
||||
else:
|
||||
scaled_avg_loss = scaler.scale(avg_loss)
|
||||
scaled_avg_loss.backward()
|
||||
scaler.minimize(optimizer, scaled_avg_loss)
|
||||
|
||||
if not isinstance(lr_scheduler, float):
|
||||
lr_scheduler.step()
|
||||
|
||||
if idx % 10 == 0:
|
||||
et = time.time()
|
||||
strs = f"epoch: [{epoch}/{EPOCH}], iter: [{idx}/{data_num}], "
|
||||
strs += f"loss: {float(avg_loss)}"
|
||||
strs += (
|
||||
f", acc_topk1: {float(acc['top1'])}, acc_top5: {float(acc['top5'])}"
|
||||
)
|
||||
strs += f", batch_time: {round(et-st, 4)} s"
|
||||
logger.info(strs)
|
||||
st = time.time()
|
||||
|
||||
if epoch % 10 == 0:
|
||||
acc = eval(config, model)
|
||||
if len(best_acc) < 1 or float(acc["top5"]) > best_acc["top5"]:
|
||||
best_acc = acc
|
||||
best_acc["epoch"] = epoch
|
||||
is_best = True
|
||||
else:
|
||||
is_best = False
|
||||
logger.info(
|
||||
f"The best acc: acc_topk1: {float(best_acc['top1'])}, acc_top5: {float(best_acc['top5'])}, best_epoch: {best_acc['epoch']}"
|
||||
)
|
||||
save_model(
|
||||
model,
|
||||
optimizer,
|
||||
config["save_model_dir"],
|
||||
logger,
|
||||
is_best,
|
||||
prefix="cls",
|
||||
)
|
||||
|
||||
|
||||
def train_distill(config, scaler=None):
|
||||
EPOCH = config["epoch"]
|
||||
topk = config["topk"]
|
||||
|
||||
batch_size = config["TRAIN"]["batch_size"]
|
||||
num_workers = config["TRAIN"]["num_workers"]
|
||||
train_loader = build_dataloader(
|
||||
"train", batch_size=batch_size, num_workers=num_workers
|
||||
)
|
||||
|
||||
# build metric
|
||||
metric_func = create_metric
|
||||
|
||||
# model = distillmv3_large_x0_5(class_dim=100)
|
||||
model = build_model(config)
|
||||
|
||||
# pact quant train
|
||||
if "quant_train" in config and config["quant_train"] is True:
|
||||
quanter = QAT(config=quant_config, act_preprocess=PACT)
|
||||
quanter.quantize(model)
|
||||
elif "prune_train" in config and config["prune_train"] is True:
|
||||
model = prune_model(model, [1, 3, 32, 32], 0.1)
|
||||
else:
|
||||
pass
|
||||
|
||||
# build_optimizer
|
||||
optimizer, lr_scheduler = create_optimizer(
|
||||
config, parameter_list=model.parameters()
|
||||
)
|
||||
|
||||
# load model
|
||||
pre_best_model_dict = load_model(config, model, optimizer)
|
||||
if len(pre_best_model_dict) > 0:
|
||||
pre_str = "The metric of loaded metric as follows {}".format(
|
||||
", ".join(["{}: {}".format(k, v) for k, v in pre_best_model_dict.items()])
|
||||
)
|
||||
logger.info(pre_str)
|
||||
|
||||
model.train()
|
||||
model = paddle.DataParallel(model)
|
||||
|
||||
# build loss function
|
||||
loss_func_distill = LossDistill(model_name_list=["student", "student1"])
|
||||
loss_func_dml = DMLLoss(model_name_pairs=["student", "student1"])
|
||||
loss_func_js = KLJSLoss(mode="js")
|
||||
|
||||
data_num = len(train_loader)
|
||||
|
||||
best_acc = {}
|
||||
for epoch in range(EPOCH):
|
||||
st = time.time()
|
||||
for idx, data in enumerate(train_loader):
|
||||
img_batch, label = data
|
||||
img_batch = paddle.transpose(img_batch, [0, 3, 1, 2])
|
||||
label = paddle.unsqueeze(label, -1)
|
||||
if scaler is not None:
|
||||
with paddle.amp.auto_cast():
|
||||
outs = model(img_batch)
|
||||
else:
|
||||
outs = model(img_batch)
|
||||
|
||||
# cal metric
|
||||
acc = metric_func(outs["student"], label)
|
||||
|
||||
# cal loss
|
||||
avg_loss = (
|
||||
loss_func_distill(outs, label)["student"]
|
||||
+ loss_func_distill(outs, label)["student1"]
|
||||
+ loss_func_dml(outs, label)["student_student1"]
|
||||
)
|
||||
|
||||
# backward
|
||||
if scaler is None:
|
||||
avg_loss.backward()
|
||||
optimizer.step()
|
||||
optimizer.clear_grad()
|
||||
else:
|
||||
scaled_avg_loss = scaler.scale(avg_loss)
|
||||
scaled_avg_loss.backward()
|
||||
scaler.minimize(optimizer, scaled_avg_loss)
|
||||
|
||||
if not isinstance(lr_scheduler, float):
|
||||
lr_scheduler.step()
|
||||
|
||||
if idx % 10 == 0:
|
||||
et = time.time()
|
||||
strs = f"epoch: [{epoch}/{EPOCH}], iter: [{idx}/{data_num}], "
|
||||
strs += f"loss: {float(avg_loss)}"
|
||||
strs += (
|
||||
f", acc_topk1: {float(acc['top1'])}, acc_top5: {float(acc['top5'])}"
|
||||
)
|
||||
strs += f", batch_time: {round(et-st, 4)} s"
|
||||
logger.info(strs)
|
||||
st = time.time()
|
||||
|
||||
if epoch % 10 == 0:
|
||||
acc = eval(config, model._layers.student)
|
||||
if len(best_acc) < 1 or float(acc["top5"]) > best_acc["top5"]:
|
||||
best_acc = acc
|
||||
best_acc["epoch"] = epoch
|
||||
is_best = True
|
||||
else:
|
||||
is_best = False
|
||||
logger.info(
|
||||
f"The best acc: acc_topk1: {float(best_acc['top1'])}, acc_top5: {float(best_acc['top5'])}, best_epoch: {best_acc['epoch']}"
|
||||
)
|
||||
|
||||
save_model(
|
||||
model,
|
||||
optimizer,
|
||||
config["save_model_dir"],
|
||||
logger,
|
||||
is_best,
|
||||
prefix="cls_distill",
|
||||
)
|
||||
|
||||
|
||||
def train_distill_multiopt(config, scaler=None):
|
||||
EPOCH = config["epoch"]
|
||||
topk = config["topk"]
|
||||
|
||||
batch_size = config["TRAIN"]["batch_size"]
|
||||
num_workers = config["TRAIN"]["num_workers"]
|
||||
train_loader = build_dataloader(
|
||||
"train", batch_size=batch_size, num_workers=num_workers
|
||||
)
|
||||
|
||||
# build metric
|
||||
metric_func = create_metric
|
||||
|
||||
# model = distillmv3_large_x0_5(class_dim=100)
|
||||
model = build_model(config)
|
||||
|
||||
# build_optimizer
|
||||
optimizer, lr_scheduler = create_optimizer(
|
||||
config, parameter_list=model.student.parameters()
|
||||
)
|
||||
optimizer1, lr_scheduler1 = create_optimizer(
|
||||
config, parameter_list=model.student1.parameters()
|
||||
)
|
||||
|
||||
# load model
|
||||
pre_best_model_dict = load_model(config, model, optimizer)
|
||||
if len(pre_best_model_dict) > 0:
|
||||
pre_str = "The metric of loaded metric as follows {}".format(
|
||||
", ".join(["{}: {}".format(k, v) for k, v in pre_best_model_dict.items()])
|
||||
)
|
||||
logger.info(pre_str)
|
||||
|
||||
# quant train
|
||||
if "quant_train" in config and config["quant_train"] is True:
|
||||
quanter = QAT(config=quant_config, act_preprocess=PACT)
|
||||
quanter.quantize(model)
|
||||
elif "prune_train" in config and config["prune_train"] is True:
|
||||
model = prune_model(model, [1, 3, 32, 32], 0.1)
|
||||
else:
|
||||
pass
|
||||
|
||||
model.train()
|
||||
|
||||
model = paddle.DataParallel(model)
|
||||
|
||||
# build loss function
|
||||
loss_func_distill = LossDistill(model_name_list=["student", "student1"])
|
||||
loss_func_dml = DMLLoss(model_name_pairs=["student", "student1"])
|
||||
loss_func_js = KLJSLoss(mode="js")
|
||||
|
||||
data_num = len(train_loader)
|
||||
best_acc = {}
|
||||
for epoch in range(EPOCH):
|
||||
st = time.time()
|
||||
for idx, data in enumerate(train_loader):
|
||||
img_batch, label = data
|
||||
img_batch = paddle.transpose(img_batch, [0, 3, 1, 2])
|
||||
label = paddle.unsqueeze(label, -1)
|
||||
|
||||
if scaler is not None:
|
||||
with paddle.amp.auto_cast():
|
||||
outs = model(img_batch)
|
||||
else:
|
||||
outs = model(img_batch)
|
||||
|
||||
# cal metric
|
||||
acc = metric_func(outs["student"], label)
|
||||
|
||||
# cal loss
|
||||
avg_loss = (
|
||||
loss_func_distill(outs, label)["student"]
|
||||
+ loss_func_dml(outs, label)["student_student1"]
|
||||
)
|
||||
avg_loss1 = (
|
||||
loss_func_distill(outs, label)["student1"]
|
||||
+ loss_func_dml(outs, label)["student_student1"]
|
||||
)
|
||||
|
||||
if scaler is None:
|
||||
# backward
|
||||
avg_loss.backward(retain_graph=True)
|
||||
optimizer.step()
|
||||
optimizer.clear_grad()
|
||||
|
||||
avg_loss1.backward()
|
||||
optimizer1.step()
|
||||
optimizer1.clear_grad()
|
||||
else:
|
||||
scaled_avg_loss = scaler.scale(avg_loss)
|
||||
scaled_avg_loss.backward()
|
||||
scaler.minimize(optimizer, scaled_avg_loss)
|
||||
|
||||
scaled_avg_loss = scaler.scale(avg_loss1)
|
||||
scaled_avg_loss.backward()
|
||||
scaler.minimize(optimizer1, scaled_avg_loss)
|
||||
|
||||
if not isinstance(lr_scheduler, float):
|
||||
lr_scheduler.step()
|
||||
if not isinstance(lr_scheduler1, float):
|
||||
lr_scheduler1.step()
|
||||
|
||||
if idx % 10 == 0:
|
||||
et = time.time()
|
||||
strs = f"epoch: [{epoch}/{EPOCH}], iter: [{idx}/{data_num}], "
|
||||
strs += f"loss: {float(avg_loss)}, loss1: {float(avg_loss1)}"
|
||||
strs += (
|
||||
f", acc_topk1: {float(acc['top1'])}, acc_top5: {float(acc['top5'])}"
|
||||
)
|
||||
strs += f", batch_time: {round(et-st, 4)} s"
|
||||
logger.info(strs)
|
||||
st = time.time()
|
||||
|
||||
if epoch % 10 == 0:
|
||||
acc = eval(config, model._layers.student)
|
||||
if len(best_acc) < 1 or float(acc["top5"]) > best_acc["top5"]:
|
||||
best_acc = acc
|
||||
best_acc["epoch"] = epoch
|
||||
is_best = True
|
||||
else:
|
||||
is_best = False
|
||||
logger.info(
|
||||
f"The best acc: acc_topk1: {float(best_acc['top1'])}, acc_top5: {float(best_acc['top5'])}, best_epoch: {best_acc['epoch']}"
|
||||
)
|
||||
save_model(
|
||||
model,
|
||||
[optimizer, optimizer1],
|
||||
config["save_model_dir"],
|
||||
logger,
|
||||
is_best,
|
||||
prefix="cls_distill_multiopt",
|
||||
)
|
||||
|
||||
|
||||
def eval(config, model):
|
||||
batch_size = config["VALID"]["batch_size"]
|
||||
num_workers = config["VALID"]["num_workers"]
|
||||
valid_loader = build_dataloader(
|
||||
"test", batch_size=batch_size, num_workers=num_workers
|
||||
)
|
||||
|
||||
# build metric
|
||||
metric_func = create_metric
|
||||
|
||||
outs = []
|
||||
labels = []
|
||||
for idx, data in enumerate(valid_loader):
|
||||
img_batch, label = data
|
||||
img_batch = paddle.transpose(img_batch, [0, 3, 1, 2])
|
||||
label = paddle.unsqueeze(label, -1)
|
||||
out = model(img_batch)
|
||||
|
||||
outs.append(out)
|
||||
labels.append(label)
|
||||
|
||||
outs = paddle.concat(outs, axis=0)
|
||||
labels = paddle.concat(labels, axis=0)
|
||||
acc = metric_func(outs, labels)
|
||||
|
||||
strs = f"The metric are as follows: acc_topk1: {float(acc['top1'])}, acc_top5: {float(acc['top5'])}"
|
||||
logger.info(strs)
|
||||
return acc
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
config, logger = preprocess(is_train=False)
|
||||
|
||||
# AMP scaler
|
||||
scaler = amp_scaler(config)
|
||||
|
||||
model_type = config["model_type"]
|
||||
|
||||
if model_type == "cls":
|
||||
train(config)
|
||||
elif model_type == "cls_distill":
|
||||
train_distill(config)
|
||||
elif model_type == "cls_distill_multiopt":
|
||||
train_distill_multiopt(config)
|
||||
else:
|
||||
raise ValueError("model_type should be one of ['']")
|
||||
Reference in New Issue
Block a user