mirror of
https://github.com/netfun2000/dddd_trainer.git
synced 2026-08-13 01:01:43 +08:00
beta0.1
This commit is contained in:
sml2h3
@@ -0,0 +1,116 @@
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
|
||||
import tqdm
|
||||
|
||||
from configs import Config
|
||||
from loguru import logger
|
||||
|
||||
|
||||
class CacheData:
|
||||
def __init__(self, project_name: str):
|
||||
self.project_name = project_name
|
||||
self.project_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects",
|
||||
project_name)
|
||||
if os.path.exists(self.project_path):
|
||||
self.cache_path = os.path.join(self.project_path, "cache")
|
||||
else:
|
||||
logger.error("Project {} is not exists!".format(project_name))
|
||||
exit()
|
||||
self.config = Config(project_name)
|
||||
self.conf = self.config.load_config()
|
||||
self.bath_path = self.conf['System']['Path']
|
||||
self.allow_ext = []
|
||||
|
||||
def cache(self, base_path: str, search_type="name"):
|
||||
self.bath_path = base_path
|
||||
self.allow_ext = self.conf["System"]["Allow_Ext"]
|
||||
if search_type == "name":
|
||||
self.__get_label_from_name(base_path=base_path)
|
||||
else:
|
||||
self.__get_label_from_file(base_path=base_path)
|
||||
|
||||
def __get_label_from_name(self, base_path: str):
|
||||
files = os.listdir(base_path)
|
||||
logger.info("\nFiles number is {}.".format(len(files)))
|
||||
self.__collect_data(files, base_path)
|
||||
|
||||
def __get_label_from_file(self, base_path: str):
|
||||
labels_path = os.path.join(base_path, "labels.txt")
|
||||
images_path = os.path.join(base_path, "images")
|
||||
if not os.path.exists(labels_path):
|
||||
logger.error("\nThe file labels.txt not found in path ----> {}".format(base_path))
|
||||
exit()
|
||||
if not os.path.exists(images_path) or not os.path.isdir(images_path):
|
||||
logger.error("\nThe dir {} not found in path ----> {}".format(images_path, base_path))
|
||||
exit()
|
||||
files = os.listdir(images_path)
|
||||
logger.info("\nFiles number is {}.".format(len(files)))
|
||||
with open(labels_path, "r", encoding="utf-8") as f:
|
||||
labels_lines = f.readlines()
|
||||
labels_lines = [line.replace("\r", "").replace("\n", "") for line in labels_lines]
|
||||
labels_filename_lines = [line.split("\t")[0] for line in labels_lines]
|
||||
logger.info("\nLabels number is {}.".format(len(labels_lines)))
|
||||
logger.info("\nChecking labels.txt ...")
|
||||
error_files = set(labels_filename_lines).difference(set(files))
|
||||
logger.info("\nCheck labels.txt end! {} errors!".format(len(error_files)))
|
||||
for ef in error_files:
|
||||
labels_lines.remove(ef)
|
||||
del files
|
||||
self.__collect_data(labels_lines, images_path, is_file=True)
|
||||
|
||||
def __collect_data(self, lines, base_path, is_file=False):
|
||||
labels = []
|
||||
caches = []
|
||||
if not self.conf['Model']['Word']:
|
||||
labels.append(" ")
|
||||
for file in tqdm.tqdm(lines):
|
||||
if is_file:
|
||||
line_list = file.split('\t')
|
||||
filename = line_list[0]
|
||||
label = line_list[1]
|
||||
else:
|
||||
filename = file
|
||||
label = "_".join(filename.split("_")[:-1])
|
||||
label = label.replace(" ", "")
|
||||
if filename.split('.')[-1] in self.allow_ext:
|
||||
if " " in filename:
|
||||
logger.warning("The {} has black. We will remove it!".format(filename))
|
||||
continue
|
||||
caches.append('\t'.join([filename, label]))
|
||||
if not self.conf['Model']['Word']:
|
||||
label = list(label)
|
||||
labels.extend(label)
|
||||
else:
|
||||
labels.append(label)
|
||||
|
||||
else:
|
||||
logger.warning("\nFile({}) has a suffix that is not allowed! We will remove it!".format(file))
|
||||
labels = list(set(labels))
|
||||
logger.info("\nCoolect labels is {}".format(json.dumps(labels, ensure_ascii=False)))
|
||||
self.conf['System']['Path'] = base_path
|
||||
self.conf['Model']['CharSet'] = labels
|
||||
self.config.make_config(config_dict=self.conf, single=self.conf['Model']['Word'])
|
||||
logger.info("\nWriting Cache Data!")
|
||||
del lines
|
||||
logger.info("\nCache Data Number is {}".format(len(caches)))
|
||||
logger.info("\nWriting Train and Val File.".format(len(caches)))
|
||||
val = self.conf['System']['Val']
|
||||
if 0 < val < 1:
|
||||
val_num = int(len(caches) * val)
|
||||
elif 1 < val < len(caches):
|
||||
val_num = int(val)
|
||||
else:
|
||||
logger.error("val setting vaild!")
|
||||
exit()
|
||||
random.shuffle(caches)
|
||||
train_set = caches[val_num:]
|
||||
val_set = caches[:val_num]
|
||||
del caches
|
||||
with open(os.path.join(self.cache_path, "cache.train.tmp"), 'w', encoding="utf-8") as f:
|
||||
f.write("\n".join(train_set))
|
||||
with open(os.path.join(self.cache_path, "cache.val.tmp"), 'w', encoding="utf-8") as f:
|
||||
f.write("\n".join(val_set))
|
||||
logger.info("\nTrain Data Number is {}".format(len(train_set)))
|
||||
logger.info("\nVal Data Number is {}".format(len(val_set)))
|
||||
@@ -0,0 +1,162 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
import torch
|
||||
import tqdm
|
||||
import numpy as np
|
||||
|
||||
from configs import Config
|
||||
from loguru import logger
|
||||
|
||||
import torchvision
|
||||
from torch.utils.data import DataLoader, Dataset, TensorDataset
|
||||
|
||||
|
||||
class LoadCache(Dataset):
|
||||
def __init__(self, cache_path: str, path: str, word: bool, image_channel: int, resize: list, charset: list):
|
||||
self.cache_path = cache_path
|
||||
self.path = path
|
||||
self.word = word
|
||||
self.ImageChannel = image_channel
|
||||
self.resize = resize
|
||||
self.charset = charset
|
||||
|
||||
logger.info("\nReading Cache File... ----> {}".format(self.cache_path))
|
||||
|
||||
with open(self.cache_path, 'r', encoding='utf-8') as f:
|
||||
caches = f.readlines()
|
||||
self.caches = []
|
||||
for cache in tqdm.tqdm(caches):
|
||||
cache = cache.replace("\r", "").replace("\n", "").split("\t")
|
||||
self.caches.append(cache)
|
||||
del caches
|
||||
|
||||
self.caches_num = len(self.caches)
|
||||
logger.info("\nRead Cache File End! Caches Num is {}.".format(self.caches_num))
|
||||
|
||||
def __len__(self):
|
||||
return self.caches_num
|
||||
|
||||
def __getitem__(self, idx):
|
||||
try:
|
||||
data = self.caches[idx]
|
||||
image_name = data[0]
|
||||
image_label = data[1]
|
||||
image_path = os.path.join(self.path, image_name)
|
||||
if not self.word:
|
||||
image_label = list(image_label)
|
||||
else:
|
||||
image_label = [image_label]
|
||||
if self.ImageChannel == 1:
|
||||
mode = torchvision.io.ImageReadMode.GRAY
|
||||
else:
|
||||
mode = torchvision.io.ImageReadMode.RGB
|
||||
image = torchvision.io.read_image(image_path, mode=mode) # shape c, h, w
|
||||
image_shape = image.shape
|
||||
image_height = image_shape[1]
|
||||
image_width = image_shape[2]
|
||||
width = self.resize[0]
|
||||
height = self.resize[1]
|
||||
if self.resize[0] == -1:
|
||||
image = torchvision.transforms.Resize((height, int(image_width * (height / image_height))))(image)
|
||||
else:
|
||||
image = torchvision.transforms.Resize((height, width))(image)
|
||||
image = torchvision.transforms.ToPILImage()(image)
|
||||
label = [int(self.charset.index(item)) for item in list(image_label)]
|
||||
return image, label
|
||||
|
||||
except Exception as e:
|
||||
logger.error("\nError: {}, File: {}".format(str(e), self.caches[idx][0]))
|
||||
return None, None
|
||||
|
||||
|
||||
class GetLoader:
|
||||
def __init__(self, project_name: str):
|
||||
self.project_name = project_name
|
||||
self.project_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects",
|
||||
project_name)
|
||||
if os.path.exists(self.project_path):
|
||||
self.cache_path = os.path.join(self.project_path, "cache")
|
||||
if os.path.exists(self.cache_path):
|
||||
self.cache_train_path = os.path.join(self.cache_path, "cache.train.tmp")
|
||||
self.cache_val_path = os.path.join(self.cache_path, "cache.val.tmp")
|
||||
|
||||
if not os.path.exists(self.cache_train_path):
|
||||
logger.error("\nCache Train File {} is not exists!".format(self.cache_train_path))
|
||||
exit()
|
||||
if not os.path.exists(self.cache_val_path):
|
||||
logger.error("\nCache Val File {} is not exists!".format(self.cache_val_path))
|
||||
exit()
|
||||
|
||||
else:
|
||||
logger.error("\nCache dir {} is not exists!".format(self.cache_path))
|
||||
exit()
|
||||
else:
|
||||
logger.error("\nProject {} is not exists!".format(project_name))
|
||||
exit()
|
||||
|
||||
self.config = Config(project_name)
|
||||
self.conf = self.config.load_config()
|
||||
|
||||
self.charset = self.conf['Model']['CharSet']
|
||||
logger.info("\nCharsets is {}".format(json.dumps(self.charset, ensure_ascii=False)))
|
||||
|
||||
self.resize = [int(self.conf['Model']['ImageWidth']), int(self.conf['Model']['ImageHeight'])]
|
||||
logger.info("\nImage Resize is {}".format(json.dumps(self.resize)))
|
||||
|
||||
self.ImageChannel = self.conf['Model']['ImageChannel']
|
||||
|
||||
self.word = self.conf['Model']['Word']
|
||||
|
||||
self.path = self.conf['System']['Path']
|
||||
|
||||
self.batch_size = self.conf['Train']['BATCH_SIZE']
|
||||
|
||||
self.val_batch_size = self.conf['Train']['TEST_BATCH_SIZE']
|
||||
|
||||
logger.info("\nImage Path is {}".format(self.path))
|
||||
|
||||
self.transform_list = []
|
||||
self.transform_list.append(torchvision.transforms.ToTensor())
|
||||
if self.ImageChannel == 1:
|
||||
self.transform_list.append(torchvision.transforms.Normalize(mean=[0.456],
|
||||
std=[0.224]))
|
||||
else:
|
||||
if self.ImageChannel != 3:
|
||||
logger.error("ImageChannel must be 1 or 3!")
|
||||
exit()
|
||||
self.transform_list.append(torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225]))
|
||||
self.transform = torchvision.transforms.Compose(self.transform_list)
|
||||
tarin_loader = LoadCache(self.cache_train_path, self.path, self.word, self.ImageChannel, self.resize, self.charset)
|
||||
val_loader = LoadCache(self.cache_val_path, self.path, self.word, self.ImageChannel, self.resize, self.charset)
|
||||
self.loaders = {
|
||||
'train': DataLoader(dataset=tarin_loader, batch_size=self.batch_size, shuffle=True, drop_last=True,
|
||||
num_workers=0, collate_fn=self.collate_to_sparse),
|
||||
'val': DataLoader(dataset=val_loader, batch_size=self.val_batch_size, shuffle=True, drop_last=True,
|
||||
num_workers=0, collate_fn=self.collate_to_sparse),
|
||||
}
|
||||
|
||||
def collate_to_sparse(self, batch):
|
||||
values = []
|
||||
images = []
|
||||
shapes = []
|
||||
max_width = 0
|
||||
for n, (img, seq) in enumerate(batch):
|
||||
if img is None or seq is None:
|
||||
continue
|
||||
if len(seq) == 0: continue
|
||||
if max_width < img.size[0]:
|
||||
max_width = img.size[0]
|
||||
values.extend(seq)
|
||||
images.append(img)
|
||||
shapes.append(len(seq))
|
||||
images_pad = []
|
||||
for img in images:
|
||||
img = torchvision.transforms.Pad((0, 0, int(max_width - img.size[0]), 0))(img)
|
||||
if self.transform is not None:
|
||||
img = self.transform(img)
|
||||
images_pad.append(img)
|
||||
images_pad = torch.stack(images_pad, dim=0)
|
||||
return [images_pad, torch.FloatTensor(values), torch.IntTensor(shapes)]
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
import os
|
||||
from configs import Config
|
||||
from loguru import logger
|
||||
|
||||
|
||||
class ProjectManager:
|
||||
|
||||
def __init__(self):
|
||||
self.base_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects")
|
||||
|
||||
def create_project(self, project_name: str, single: bool = False):
|
||||
project_base_path = os.path.join(self.base_path, project_name)
|
||||
logger.info("Creating Directory... ----> {}".format(project_base_path))
|
||||
if not os.path.exists(project_base_path):
|
||||
os.mkdir(project_base_path)
|
||||
if not os.path.exists(project_base_path):
|
||||
logger.error("Directory create failed! ----> {}".format(project_base_path))
|
||||
return False
|
||||
models_path = os.path.join(project_base_path, "models")
|
||||
logger.info("Creating Directory... ----> {}".format(models_path))
|
||||
os.mkdir(models_path)
|
||||
|
||||
cache_path = os.path.join(project_base_path, "cache")
|
||||
logger.info("Creating Directory... ----> {}".format(cache_path))
|
||||
os.mkdir(cache_path)
|
||||
|
||||
checkpoints_path = os.path.join(project_base_path, "checkpoints")
|
||||
logger.info("Creating Directory... ----> {}".format(checkpoints_path))
|
||||
os.mkdir(checkpoints_path)
|
||||
|
||||
config_path = os.path.join(os.path.join(project_base_path, "config.yaml"))
|
||||
logger.info("Creating {} Config File... ----> {}".format("CNN" if single else "CRNN", config_path))
|
||||
conf = Config(project_name)
|
||||
conf.make_config(single=single)
|
||||
|
||||
logger.info("Create Project Success! ----> {}".format(project_name))
|
||||
else:
|
||||
logger.error("Directory already exists! ----> {}".format(project_base_path))
|
||||
return False
|
||||
+128
@@ -0,0 +1,128 @@
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
|
||||
import tqdm
|
||||
|
||||
from configs import Config
|
||||
from loguru import logger
|
||||
from utils import load_cache
|
||||
from nets import Net
|
||||
|
||||
|
||||
class Train:
|
||||
def __init__(self, project_name: str):
|
||||
self.project_name = project_name
|
||||
self.project_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects",
|
||||
project_name)
|
||||
self.checkpoints_path = os.path.join(self.project_path, "checkpoints")
|
||||
self.models_path = os.path.join(self.project_path, "models")
|
||||
self.config = Config(project_name)
|
||||
self.conf = self.config.load_config()
|
||||
|
||||
self.test_step = self.conf['Train']['TEST_STEP']
|
||||
self.target = self.conf['Train']['TARGET']
|
||||
self.target_acc = self.target['Accuracy']
|
||||
self.min_epoch = self.target['Epoch']
|
||||
self.max_loss = self.target['Cost']
|
||||
logger.info("\nTaget:\nmin_Accuracy: {}\nmin_Epoch: {}\nmax_Loss: {}".format(self.target_acc, self.min_epoch,
|
||||
self.max_loss))
|
||||
|
||||
logger.info("\nBuilding Net...")
|
||||
self.net = Net(self.conf)
|
||||
logger.info(self.net)
|
||||
logger.info("\nBuilding End")
|
||||
|
||||
self.use_gpu = self.conf['System']['GPU']
|
||||
if self.use_gpu:
|
||||
self.gpu_id = self.conf['System']['GPU_ID']
|
||||
logger.info("\nUSE GPU ----> {}".format(self.gpu_id))
|
||||
self.device = self.net.get_device(self.gpu_id)
|
||||
self.net.to(self.device)
|
||||
else:
|
||||
self.gpu_id = -1
|
||||
self.device = self.net.get_device(self.gpu_id)
|
||||
logger.info("\nUSE CPU".format(self.gpu_id))
|
||||
logger.info("\nGet Data Loader...")
|
||||
loaders = load_cache.GetLoader(project_name)
|
||||
self.train = loaders.loaders['train']
|
||||
self.val = loaders.loaders['val']
|
||||
logger.info("\nGet Data Loader End!")
|
||||
|
||||
self.epoch = 0
|
||||
self.step = 0
|
||||
self.loss = 0
|
||||
self.avg_loss = 0
|
||||
self.start_time = time.time()
|
||||
self.now_time = time.time()
|
||||
|
||||
def start(self):
|
||||
val_iter = iter(self.val)
|
||||
while True:
|
||||
for idx, (inputs, labels, labels_length) in enumerate(self.train):
|
||||
self.now_time = time.time()
|
||||
inputs = self.net.variable_to_device(inputs, device=self.device)
|
||||
|
||||
loss, lr = self.net.training(inputs, labels, labels_length)
|
||||
|
||||
self.avg_loss += loss
|
||||
|
||||
self.step += 1
|
||||
|
||||
if self.step % 100 == 0 and self.step % self.test_step != 0:
|
||||
logger.info("{}\tEpoch: {}\tStep: {}\tLastLoss: {}\tAvgLoss: {}\tLr: {}".format(
|
||||
time.strftime("[%Y-%m-%d-%H_%M_%S]", time.localtime(self.now_time)), self.epoch, self.step,
|
||||
str(loss), str(self.avg_loss / 100), lr
|
||||
))
|
||||
self.avg_loss = 0
|
||||
if self.step % 2000 == 0 and self.step != 0:
|
||||
model_path = os.path.join(self.checkpoints_path, "checkpoint_{}_{}_{}.tar".format(
|
||||
self.project_name, self.epoch, self.step,
|
||||
))
|
||||
self.net.save_model(model_path,
|
||||
{"net": self.net.state_dict(), "optimizer": self.net.optimizer.state_dict(),
|
||||
"epoch": self.epoch, "step": self.step})
|
||||
if self.step % self.test_step == 0:
|
||||
try:
|
||||
test_inputs, test_labels, test_labels_length = next(val_iter)
|
||||
except Exception:
|
||||
del val_iter
|
||||
val_iter = iter(self.val)
|
||||
test_inputs, test_labels, test_labels_length = next(val_iter)
|
||||
if test_inputs.shape[0] < 5:
|
||||
continue
|
||||
test_inputs = self.net.variable_to_device(test_inputs, self.device)
|
||||
self.net = self.net.train(False)
|
||||
pred_labels, labels_list, correct_list, error_list = self.net.test_op(test_inputs, test_labels,
|
||||
test_labels_length)
|
||||
self.net = self.net.train()
|
||||
accuracy = len(correct_list) / test_inputs.shape[0]
|
||||
logger.info("{}\tEpoch: {}\tStep: {}\tLastLoss: {}\tAvgLoss: {}\tLr: {}\tAcc: {}".format(
|
||||
time.strftime("[%Y-%m-%d-%H_%M_%S]", time.localtime(self.now_time)), self.epoch, self.step,
|
||||
str(loss), str(self.avg_loss / 100), lr, accuracy
|
||||
))
|
||||
if accuracy > self.target_acc and self.epoch > self.min_epoch and self.avg_loss < self.max_loss:
|
||||
logger.info("\nTraining Finished!Exporting Model...")
|
||||
dummy_input = self.net.get_random_tensor()
|
||||
input_names = ["input1"]
|
||||
output_names = ["output"]
|
||||
|
||||
if self.net.backbone.startswith("effnet"):
|
||||
self.net.cnn.set_swish(memory_efficient=False)
|
||||
self.net = self.net.eval().cpu()
|
||||
dynamic_ax = {'input1': {3: 'image_wdith'}, "output": {1: 'seq'}}
|
||||
self.net.export_onnx(self.net, dummy_input,
|
||||
os.path.join(self.models_path, "{}_{}_{}_{}_{}.onnx".format(
|
||||
self.project_name, str(accuracy), self.epoch, self.step,
|
||||
time.localtime(self.now_time)))
|
||||
, input_names, output_names, dynamic_ax)
|
||||
logger.info("\nExport Finished!Using Time: {}min".format(str(int(int(self.now_time * 1000) - int(self.start_time * 1000)) / 60)))
|
||||
exit()
|
||||
|
||||
self.epoch += 1
|
||||
self.net.scheduler.step(self.epoch)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
Train("test1")
|
||||
Reference in New Issue
Block a user