def main():
eval_batch_size = config["eval_batch_size"]
cpu_workers = config["cpu_workers"]
reload_checkpoint = config["reload_checkpoint"]
img_height = config["img_height"]
img_width = config["img_width"]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"device: {device}")
test_dataset = Synth90kDataset(
root_dir=config["data_dir"],
mode="test",
img_height=img_height,
img_width=img_width,
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=eval_batch_size,
shuffle=False,
num_workers=cpu_workers,
collate_fn=synth90k_collate_fn,
)
num_class = len(Synth90kDataset.LABEL2CHAR) + 1
crnn = CRNN(
1,
img_height,
img_width,
num_class,
map_to_seq_hidden=config["map_to_seq_hidden"],
rnn_hidden=config["rnn_hidden"],
leaky_relu=config["leaky_relu"],
)
crnn.load_state_dict(torch.load(reload_checkpoint, map_location=device))
crnn.to(device)
criterion = CTCLoss(reduction="sum")
criterion.to(device)
evaluation = evaluate(
crnn,
test_loader,
criterion,
decode_method=config["decode_method"],
beam_size=config["beam_size"],
)
print("test_evaluation: loss={loss}, acc={acc}".format(**evaluation))
def main():
eval_batch_size = config['eval_batch_size']
cpu_workers = config['cpu_workers']
reload_checkpoint = config['reload_checkpoint']
img_height = config['img_height']
img_width = config['img_width']
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'device: {device}')
test_dataset = Synth90kDataset(root_dir=config['data_dir'],
mode='test',
img_height=img_height,
img_width=img_width)
test_loader = DataLoader(dataset=test_dataset,
batch_size=eval_batch_size,
shuffle=False,
num_workers=cpu_workers,
collate_fn=synth90k_collate_fn)
num_class = len(Synth90kDataset.LABEL2CHAR) + 1
crnn = CRNN(1,
img_height,
img_width,
num_class,
map_to_seq_hidden=config['map_to_seq_hidden'],
rnn_hidden=config['rnn_hidden'],
leaky_relu=config['leaky_relu'])
crnn.load_state_dict(torch.load(reload_checkpoint, map_location=device))
crnn.to(device)
criterion = CTCLoss(reduction='sum')
criterion.to(device)
evaluation = evaluate(crnn,
test_loader,
criterion,
decode_method=config['decode_method'],
beam_size=config['beam_size'])
print('test_evaluation: loss={loss}, acc={acc}'.format(**evaluation))
def train(num_epochs, model, device, train_loader, val_loader, images, texts,
lengths, converter, optimizer, lr_scheduler, prediction_dir,
print_iter):
criterion = CTCLoss()
criterion.to(device)
images = images.to(device)
model.to(device)
for epoch in range(num_epochs):
print(epoch)
count = 0
model.train()
for i, datas in enumerate(train_loader):
datas, targets = datas
batch_size = datas.size(0)
count += batch_size
dataloader.loadData(images, datas)
t, l = converter.encode(targets)
dataloader.loadData(texts, t)
dataloader.loadData(lengths, l)
preds = model(images)
preds_size = Variable(torch.IntTensor([preds.size(0)] *
batch_size))
cost = criterion(preds, texts, preds_size, lengths) / batch_size
model.zero_grad()
cost.backward()
optimizer.step()
if count % print_iter < train_loader.batch_size:
print('epoch {} [{}/{}]loss : {}'.format(
epoch, count, len(train_loader.dataset), cost))
validation(model, device, val_loader, images, texts, lengths,
converter, prediction_dir)
save_model('{}'.format(epoch), model, optimizer, lr_scheduler)
lr_scheduler.step()
def main():
epochs = config['epochs']
train_batch_size = config['train_batch_size']
eval_batch_size = config['eval_batch_size']
lr = config['lr']
show_interval = config['show_interval']
valid_interval = config['valid_interval']
save_interval = config['save_interval']
cpu_workers = config['cpu_workers']
reload_checkpoint = config['reload_checkpoint']
valid_max_iter = config['valid_max_iter']
img_width = config['img_width']
img_height = config['img_height']
data_dir = config['data_dir']
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'device: {device}')
train_dataset = Synth90kDataset(root_dir=data_dir,
mode='train',
img_height=img_height,
img_width=img_width)
valid_dataset = Synth90kDataset(root_dir=data_dir,
mode='dev',
img_height=img_height,
img_width=img_width)
train_loader = DataLoader(dataset=train_dataset,
batch_size=train_batch_size,
shuffle=True,
num_workers=cpu_workers,
collate_fn=synth90k_collate_fn)
valid_loader = DataLoader(dataset=valid_dataset,
batch_size=eval_batch_size,
shuffle=True,
num_workers=cpu_workers,
collate_fn=synth90k_collate_fn)
num_class = len(Synth90kDataset.LABEL2CHAR) + 1
crnn = CRNN(1,
img_height,
img_width,
num_class,
map_to_seq_hidden=config['map_to_seq_hidden'],
rnn_hidden=config['rnn_hidden'],
leaky_relu=config['leaky_relu'])
if reload_checkpoint:
crnn.load_state_dict(torch.load(reload_checkpoint,
map_location=device))
crnn.to(device)
optimizer = optim.RMSprop(crnn.parameters(), lr=lr)
criterion = CTCLoss(reduction='sum')
criterion.to(device)
assert save_interval % valid_interval == 0
i = 1
for epoch in range(1, epochs + 1):
print(f'epoch: {epoch}')
tot_train_loss = 0.
tot_train_count = 0
for train_data in train_loader:
loss = train_batch(crnn, train_data, optimizer, criterion, device)
train_size = train_data[0].size(0)
tot_train_loss += loss
tot_train_count += train_size
if i % show_interval == 0:
print('train_batch_loss[', i, ']: ', loss / train_size)
if i % valid_interval == 0:
evaluation = evaluate(crnn,
valid_loader,
criterion,
decode_method=config['decode_method'],
beam_size=config['beam_size'])
print('valid_evaluation: loss={loss}, acc={acc}'.format(
**evaluation))
if i % save_interval == 0:
prefix = 'crnn'
loss = evaluation['loss']
save_model_path = os.path.join(
config['checkpoints_dir'],
f'{prefix}_{i:06}_loss{loss}.pt')
torch.save(crnn.state_dict(), save_model_path)
print('save model at ', save_model_path)
i += 1
print('train_loss: ', tot_train_loss / tot_train_count)
def main():
epochs = config["epochs"]
train_batch_size = config["train_batch_size"]
eval_batch_size = config["eval_batch_size"]
lr = config["lr"]
show_interval = config["show_interval"]
valid_interval = config["valid_interval"]
save_interval = config["save_interval"]
cpu_workers = config["cpu_workers"]
reload_checkpoint = config["reload_checkpoint"]
valid_max_iter = config["valid_max_iter"]
img_width = config["img_width"]
img_height = config["img_height"]
data_dir = config["data_dir"]
torch.cuda.empty_cache()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"device: {device}")
# # Extracts metadata related to the file path, label, and image width + height.
# wbsin_dir = Path().cwd() / "data" / "processed" / "cropped_wbsin"
# wbsin_meta_df = extract_jpg_meta(img_dir=wbsin_dir, img_type="wbsin")
# # Saves the extracted metadata.
# interim_path = Path.cwd() / "data" / "interim"
# interim_path.mkdir(parents=True, exist_ok=True)
# wbsin_meta_df.to_csv(interim_path / "wbsin_meta.csv", index=False)
X_transforms = transforms.Compose([
transforms.Grayscale(num_output_channels=1),
transforms.Resize((160, 1440)),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, )),
])
wbsin_dataset = WbsinImageDataset(
meta_file=(Path.cwd() / "data" / "processed" /
"processed_wbsin_meta.csv"),
transform=X_transforms,
)
train_size = int(0.8 * len(wbsin_dataset))
test_size = len(wbsin_dataset) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(
wbsin_dataset,
[train_size, test_size],
generator=torch.Generator().manual_seed(42),
)
# Save the test_dataset
train_dataloader = DataLoader(
train_dataset,
batch_size=train_batch_size,
shuffle=True,
num_workers=cpu_workers,
collate_fn=wbsin_collate_fn,
)
test_dataloader = DataLoader(
test_dataset,
batch_size=eval_batch_size,
shuffle=True,
num_workers=cpu_workers,
collate_fn=wbsin_collate_fn,
)
num_class = len(WbsinImageDataset.LABEL2CHAR) + 1
crnn = CRNN(
1,
img_height,
img_width,
num_class,
map_to_seq_hidden=config["map_to_seq_hidden"],
rnn_hidden=config["rnn_hidden"],
leaky_relu=config["leaky_relu"],
)
if reload_checkpoint:
crnn.load_state_dict(torch.load(reload_checkpoint,
map_location=device))
crnn.to(device)
optimizer = optim.RMSprop(crnn.parameters(), lr=lr)
criterion = CTCLoss(reduction="sum")
criterion.to(device)
assert save_interval % valid_interval == 0
i = 1
for epoch in range(1, epochs + 1):
print(f"epoch: {epoch}")
tot_train_loss = 0.0
tot_train_count = 0
for train_data in train_dataloader:
loss = train_batch(crnn, train_data, optimizer, criterion, device)
train_size = train_data[0].size(0)
tot_train_loss += loss
tot_train_count += train_size
if i % show_interval == 0:
print("train_batch_loss[", i, "]: ", loss / train_size)
if i % valid_interval == 0:
evaluation = evaluate(
crnn,
test_dataloader,
criterion,
decode_method=config["decode_method"],
beam_size=config["beam_size"],
)
print(
"valid_evaluation: loss={loss}, acc={acc}, char_acc={char_acc}"
.format(**evaluation))
if i % save_interval == 0:
prefix = "crnn"
loss = evaluation["loss"]
save_model_path = os.path.join(
config["checkpoints_dir"],
f"{prefix}_{i:06}_loss{loss}.pt")
torch.save(crnn.state_dict(), save_model_path)
print("save model at ", save_model_path)
i += 1
print("train_loss: ", tot_train_loss / tot_train_count)
def train(args):
init_distributed_mode(args)
print(args)
device = torch.device('cuda:{}'.format(args.local_rank) if args.device ==
'cuda' and torch.cuda.is_available() else 'cpu')
torch.backends.cudnn.benchmark = True
# data loader
data_set = Generator(cfg.word.get_all_words(), args.direction)
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
data_set)
else:
train_sampler = torch.utils.data.RandomSampler(data_set)
data_loader = DataLoader(data_set,
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=args.workers)
# model
model = crnn.CRNN(len(data_set.alpha))
model = model.to(device)
criterion = CTCLoss()
criterion = criterion.to(device)
if args.distributed and args.sync_bn:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
# optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
optimizer = optim.Adadelta(model.parameters(),
weight_decay=args.weight_decay)
# lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
model_without_ddp = model.module
# 加载预训练模型
if args.init_epoch > 0:
checkpoint = torch.load(os.path.join(
args.output_dir,
'crnn.{}.{:03d}.pth'.format(args.direction, args.init_epoch)),
map_location='cpu')
optimizer.load_state_dict(checkpoint['optimizer'])
# lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
model_without_ddp.load_state_dict(checkpoint['model'])
# log
writer = SummaryWriter(
log_dir=cfg.log_dir) if utils.is_main_process() else None
# train
model.train()
for epoch in range(args.init_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
# 训练
loss = train_one_epoch(model, criterion, optimizer, data_loader,
device, epoch, args)
# 记录日志
utils.add_scalar_on_master(writer, 'scalar/lr',
optimizer.param_groups[0]['lr'], epoch + 1)
utils.add_scalar_on_master(writer, 'scalar/train_loss', loss,
epoch + 1)
utils.add_weight_history_on_master(writer, model_without_ddp,
epoch + 1)
# 更新lr
# lr_scheduler.step(epoch)
# 保存模型
if args.output_dir:
checkpoint = {
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
# 'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch + 1,
'args': args
}
utils.save_on_master(
checkpoint,
os.path.join(
args.output_dir,
'crnn.{}.{:03d}.pth'.format(args.direction, epoch + 1)))
if utils.is_main_process():
writer.close()
def train(num_epochs, model, device, train_loader, val_loader, images, texts, lengths, converter, optimizer, lr_scheduler, prediction_dir, print_iter) :
criterion = CTCLoss()
criterion.to(device)
images = images.to(device)
model.to(device)
for epoch in range(num_epochs) :
print(epoch)
count = 0
model.train()
for i, datas in enumerate(train_loader) :
datas, targets = datas
batch_size = datas.size(0)
count += batch_size
dataloader.loadData(images, datas)
t, l = converter.encode(targets)
dataloader.loadData(texts, t)
dataloader.loadData(lengths, l)
preds = model(images)
print('preds size : ', preds.size())
preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
print({'preds_size size : ', preds_size.size())
cost = criterion(preds, texts, preds_size, lengths) / batch_size
model.zero_grad()
cost.backward()
optimizer.step()
if count % print_iter < train_loader.batch_size :
print('epoch {} [{}/{}]loss : {}'.format(epoch, count, len(train_loader.dataset), cost))
validation(model, device, val_loader, images, texts, lengths, converter, prediction_dir)
save_model('{}'.format(epoch), model, optimizer, lr_scheduler)
lr_scheduler.step()
def test(model, device, test_loader, images, texts, lengths, converter, prediction_dir) :
model.to(device)
images = images.to(device)
model.eval()
pred_json = {}
pred_list = []
make_folder(prediction_dir)
for i, datas in enumerate(test_loader) :
datas, targets = datas
batch_size = datas.size(0)
dataloader.loadData(images, datas)
t, l = converter.encode(targets)
dataloader.loadData(texts, t)
dataloader.loadData(lengths, l)
preds = model(images)
preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
_, preds = preds.max(2)
preds = preds.transpose(1, 0).contiguous().view(-1)
pred_string = converter.decode(preds.data, preds_size.data, raw=False)
print('per batch pred_string : ',pred_string)
pred_dict = {'image_path' : test_loader.dataset.get_img_path(i), 'prediction' : pred_string}
pred_list.append(pred_dict)
pred_json = {'predict' : pred_list}
with open(os.path.join(prediction_dir, 'predict.json'), 'w') as save_json :
json.dump(pred_json, save_json, indent=2, ensure_ascii=False)
def validation(model, device, val_loader, images, texts, lengths, converter, prediction_dir) :
test(model, device, val_loader, images, texts, lengths, converter, prediction_dir)
print('validation test finish')
DR_PATH = os.path.join(prediction_dir, 'predict.json')
GT_PATH = glob.glob(val_loader.dataset.get_root() + '/*.json')[0]
res = evaluation.evaluation_metrics(DR_PATH, GT_PATH)
print('validation : ', res)
try:
from nipa import nipa_data
DATASET_PATH = nipa_data.get_data_root('deepfake')
except:
DATASET_PATH = os.path.join('../data')
def main() :
# mode argument
args = argparse.ArgumentParser()
args.add_argument("--letter", type=str, default=" ,.()\'\"?!01234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ가각간갇갈갉갊감갑값갓갔강갖갗같갚갛개객갠갤갬갭갯갰갱갸갹갼걀걋걍걔걘걜거걱건걷걸걺검겁것겄겅겆겉겊겋게겐겔겜겝겟겠겡겨격겪견겯결겸겹겻겼경곁계곈곌곕곗고곡곤곧골곪곬곯곰곱곳공곶과곽관괄괆괌괍괏광괘괜괠괩괬괭괴괵괸괼굄굅굇굉교굔굘굡굣구국군굳굴굵굶굻굼굽굿궁궂궈궉권궐궜궝궤궷귀귁귄귈귐귑귓규균귤그극근귿글긁금급긋긍긔기긱긴긷길긺김깁깃깅깆깊까깍깎깐깔깖깜깝깟깠깡깥깨깩깬깰깸깹깻깼깽꺄꺅꺌꺼꺽꺾껀껄껌껍껏껐껑께껙껜껨껫껭껴껸껼꼇꼈꼍꼐꼬꼭꼰꼲꼴꼼꼽꼿꽁꽂꽃꽈꽉꽐꽜꽝꽤꽥꽹꾀꾄꾈꾐꾑꾕꾜꾸꾹꾼꿀꿇꿈꿉꿋꿍꿎꿔꿜꿨꿩꿰꿱꿴꿸뀀뀁뀄뀌뀐뀔뀜뀝뀨끄끅끈끊끌끎끓끔끕끗끙끝끼끽낀낄낌낍낏낑나낙낚난낟날낡낢남납낫났낭낮낯낱낳내낵낸낼냄냅냇냈냉냐냑냔냘냠냥너넉넋넌널넒넓넘넙넛넜넝넣네넥넨넬넴넵넷넸넹녀녁년녈념녑녔녕녘녜녠노녹논놀놂놈놉놋농높놓놔놘놜놨뇌뇐뇔뇜뇝뇟뇨뇩뇬뇰뇹뇻뇽누눅눈눋눌눔눕눗눙눠눴눼뉘뉜뉠뉨뉩뉴뉵뉼늄늅늉느늑는늘늙늚늠늡늣능늦늪늬늰늴니닉닌닐닒님닙닛닝닢다닥닦단닫달닭닮닯닳담답닷닸당닺닻닿대댁댄댈댐댑댓댔댕댜더덕덖던덛덜덞덟덤덥덧덩덫덮데덱덴델뎀뎁뎃뎄뎅뎌뎐뎔뎠뎡뎨뎬도독돈돋돌돎돐돔돕돗동돛돝돠돤돨돼됐되된될됨됩됫됴두둑둔둘둠둡둣둥둬뒀뒈뒝뒤뒨뒬뒵뒷뒹듀듄듈듐듕드득든듣들듦듬듭듯등듸디딕딘딛딜딤딥딧딨딩딪따딱딴딸땀땁땃땄땅땋때땍땐땔땜땝땟땠땡떠떡떤떨떪떫떰떱떳떴떵떻떼떽뗀뗄뗌뗍뗏뗐뗑뗘뗬또똑똔똘똥똬똴뙈뙤뙨뚜뚝뚠뚤뚫뚬뚱뛔뛰뛴뛸뜀뜁뜅뜨뜩뜬뜯뜰뜸뜹뜻띄띈띌띔띕띠띤띨띰띱띳띵라락란랄람랍랏랐랑랒랖랗래랙랜랠램랩랫랬랭랴략랸럇량러럭런럴럼럽럿렀렁렇레렉렌렐렘렙렛렝려력련렬렴렵렷렸령례롄롑롓로록론롤롬롭롯롱롸롼뢍뢨뢰뢴뢸룀룁룃룅료룐룔룝룟룡루룩룬룰룸룹룻룽뤄뤘뤠뤼뤽륀륄륌륏륑류륙륜률륨륩륫륭르륵른를름릅릇릉릊릍릎리릭린릴림립릿링마막만많맏말맑맒맘맙맛망맞맡맣매맥맨맬맴맵맷맸맹맺먀먁먈먕머먹먼멀멂멈멉멋멍멎멓메멕멘멜멤멥멧멨멩며멱면멸몃몄명몇몌모목몫몬몰몲몸몹못몽뫄뫈뫘뫙뫼묀묄묍묏묑묘묜묠묩묫무묵묶문묻물묽묾뭄뭅뭇뭉뭍뭏뭐뭔뭘뭡뭣뭬뮈뮌뮐뮤뮨뮬뮴뮷므믄믈믐믓미믹민믿밀밂밈밉밋밌밍및밑바박밖밗반받발밝밞밟밤밥밧방밭배백밴밸뱀뱁뱃뱄뱅뱉뱌뱍뱐뱝버벅번벋벌벎범법벗벙벚베벡벤벧벨벰벱벳벴벵벼벽변별볍볏볐병볕볘볜보복볶본볼봄봅봇봉봐봔봤봬뵀뵈뵉뵌뵐뵘뵙뵤뵨부북분붇불붉붊붐붑붓붕붙붚붜붤붰붸뷔뷕뷘뷜뷩뷰뷴뷸븀븃븅브븍븐블븜븝븟비빅빈빌빎빔빕빗빙빚빛빠빡빤빨빪빰빱빳빴빵빻빼빽뺀뺄뺌뺍뺏뺐뺑뺘뺙뺨뻐뻑뻔뻗뻘뻠뻣뻤뻥뻬뼁뼈뼉뼘뼙뼛뼜뼝뽀뽁뽄뽈뽐뽑뽕뾔뾰뿅뿌뿍뿐뿔뿜뿟뿡쀼쁑쁘쁜쁠쁨쁩삐삑삔삘삠삡삣삥사삭삯산삳살삵삶삼삽삿샀상샅새색샌샐샘샙샛샜생샤샥샨샬샴샵샷샹섀섄섈섐섕서석섞섟선섣설섦섧섬섭섯섰성섶세섹센셀셈셉셋셌셍셔셕션셜셤셥셧셨셩셰셴셸솅소속솎손솔솖솜솝솟송솥솨솩솬솰솽쇄쇈쇌쇔쇗쇘쇠쇤쇨쇰쇱쇳쇼쇽숀숄숌숍숏숑수숙순숟술숨숩숫숭숯숱숲숴쉈쉐쉑쉔쉘쉠쉥쉬쉭쉰쉴쉼쉽쉿슁슈슉슐슘슛슝스슥슨슬슭슴습슷승시식신싣실싫심십싯싱싶싸싹싻싼쌀쌈쌉쌌쌍쌓쌔쌕쌘쌜쌤쌥쌨쌩썅써썩썬썰썲썸썹썼썽쎄쎈쎌쏀쏘쏙쏜쏟쏠쏢쏨쏩쏭쏴쏵쏸쐈쐐쐤쐬쐰쐴쐼쐽쑈쑤쑥쑨쑬쑴쑵쑹쒀쒔쒜쒸쒼쓩쓰쓱쓴쓸쓺쓿씀씁씌씐씔씜씨씩씬씰씸씹씻씽아악안앉않알앍앎앓암압앗았앙앝앞애액앤앨앰앱앳앴앵야약얀얄얇얌얍얏양얕얗얘얜얠얩어억언얹얻얼얽얾엄업없엇었엉엊엌엎에엑엔엘엠엡엣엥여역엮연열엶엷염엽엾엿였영옅옆옇예옌옐옘옙옛옜오옥온올옭옮옰옳옴옵옷옹옻와왁완왈왐왑왓왔왕왜왝왠왬왯왱외왹왼욀욈욉욋욍요욕욘욜욤욥욧용우욱운울욹욺움웁웃웅워웍원월웜웝웠웡웨웩웬웰웸웹웽위윅윈윌윔윕윗윙유육윤율윰윱윳융윷으윽은을읊음읍읏응읒읓읔읕읖읗의읩읜읠읨읫이익인일읽읾잃임입잇있잉잊잎자작잔잖잗잘잚잠잡잣잤장잦재잭잰잴잼잽잿쟀쟁쟈쟉쟌쟎쟐쟘쟝쟤쟨쟬저적전절젊점접젓정젖제젝젠젤젬젭젯젱져젼졀졈졉졌졍졔조족존졸졺좀좁좃종좆좇좋좌좍좔좝좟좡좨좼좽죄죈죌죔죕죗죙죠죡죤죵주죽준줄줅줆줌줍줏중줘줬줴쥐쥑쥔쥘쥠쥡쥣쥬쥰쥴쥼즈즉즌즐즘즙즛증지직진짇질짊짐집짓징짖짙짚짜짝짠짢짤짧짬짭짯짰짱째짹짼쨀쨈쨉쨋쨌쨍쨔쨘쨩쩌쩍쩐쩔쩜쩝쩟쩠쩡쩨쩽쪄쪘쪼쪽쫀쫄쫌쫍쫏쫑쫓쫘쫙쫠쫬쫴쬈쬐쬔쬘쬠쬡쭁쭈쭉쭌쭐쭘쭙쭝쭤쭸쭹쮜쮸쯔쯤쯧쯩찌찍찐찔찜찝찡찢찧차착찬찮찰참찹찻찼창찾채책챈챌챔챕챗챘챙챠챤챦챨챰챵처척천철첨첩첫첬청체첵첸첼쳄쳅쳇쳉쳐쳔쳤쳬쳰촁초촉촌촐촘촙촛총촤촨촬촹최쵠쵤쵬쵭쵯쵱쵸춈추축춘출춤춥춧충춰췄췌췐취췬췰췸췹췻췽츄츈츌츔츙츠측츤츨츰츱츳층치칙친칟칠칡침칩칫칭카칵칸칼캄캅캇캉캐캑캔캘캠캡캣캤캥캬캭컁커컥컨컫컬컴컵컷컸컹케켁켄켈켐켑켓켕켜켠켤켬켭켯켰켱켸코콕콘콜콤콥콧콩콰콱콴콸쾀쾅쾌쾡쾨쾰쿄쿠쿡쿤쿨쿰쿱쿳쿵쿼퀀퀄퀑퀘퀭퀴퀵퀸퀼큄큅큇큉큐큔큘큠크큭큰클큼큽킁키킥킨킬킴킵킷킹타탁탄탈탉탐탑탓탔탕태택탠탤탬탭탯탰탱탸턍터턱턴털턺텀텁텃텄텅테텍텐텔템텝텟텡텨텬텼톄톈토톡톤톨톰톱톳통톺톼퇀퇘퇴퇸툇툉툐투툭툰툴툼툽툿퉁퉈퉜퉤튀튁튄튈튐튑튕튜튠튤튬튱트특튼튿틀틂틈틉틋틔틘틜틤틥티틱틴틸팀팁팃팅파팍팎판팔팖팜팝팟팠팡팥패팩팬팰팸팹팻팼팽퍄퍅퍼퍽펀펄펌펍펏펐펑페펙펜펠펨펩펫펭펴편펼폄폅폈평폐폘폡폣포폭폰폴폼폽폿퐁퐈퐝푀푄표푠푤푭푯푸푹푼푿풀풂품풉풋풍풔풩퓌퓐퓔퓜퓟퓨퓬퓰퓸퓻퓽프픈플픔픕픗피픽핀필핌핍핏핑하학한할핥함합핫항해핵핸핼햄햅햇했행햐향허헉헌헐헒험헙헛헝헤헥헨헬헴헵헷헹혀혁현혈혐협혓혔형혜혠혤혭호혹혼홀홅홈홉홋홍홑화확환활홧황홰홱홴횃횅회획횐횔횝횟횡효횬횰횹횻후훅훈훌훑훔훗훙훠훤훨훰훵훼훽휀휄휑휘휙휜휠휨휩휫휭휴휵휸휼흄흇흉흐흑흔흖흗흘흙흠흡흣흥흩희흰흴흼흽힁히힉힌힐힘힙힛힝")
args.add_argument("--lr", type=float, default=0.0001)
args.add_argument("--cuda", type=bool, default=True)
args.add_argument("--num_epochs", type=int, default=50000)
args.add_argument("--model_name", type=str, default="")
args.add_argument("--batch", type=int, default=1)
args.add_argument("--mode", type=str, default="train")
args.add_argument("--prediction_dir", type=str, default="prediction")
args.add_argument("--print_iter", type=int, default=10)
config = args.parse_args()
letter = config.letter
lr = config.lr
cuda = config.cuda
num_epochs = config.num_epochs
model_name = config.model_name
batch = config.batch
mode = config.mode
prediction_dir = config.prediction_dir
print_iter = config.print_iter
imgH = 32
imgW = 200
nclass = len(letter) + 1
nc = 1
new_model = model.CRNN(imgH, nc, nclass, 256)
new_model.apply(model.weights_init)
device = torch.device('cuda') if cuda else torch.device('cpu')
converter = dataloader.strLabelConverter(letter)
images = torch.FloatTensor(batch, 1, imgH, imgW)
texts = torch.IntTensor(batch * 1000)
lengths = torch.IntTensor(batch)
images = Variable(images)
texts = Variable(texts)
lengths = Variable(lengths)
#check parameter of model
print("------------------------------------------------------------")
total_params = sum(p.numel() for p in new_model.parameters())
print("num of parameter : ",total_params)
trainable_params = sum(p.numel() for p in new_model.parameters() if p.requires_grad)
print("num of trainable_ parameter :",trainable_params)
print("------------------------------------------------------------")
if mode == 'train' :
print('trian start')
train_loader = data_loader(
DATASET_PATH, batch, imgH, imgW, phase='train')
val_loader = data_loader(
DATASET_PATH, batch, imgH, imgW, phase='val')
params = [p for p in new_model.parameters() if p.requires_grad]
optimizer = optim.Adam(params, lr=lr, betas=(0.5, 0.999))
lr_scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=40,gamma=0.1)
train(num_epochs, new_model, device, train_loader, val_loader, images, texts, lengths, converter, optimizer, lr_scheduler, prediction_dir, print_iter)
elif mode == 'test' :
print('test start')
test_loader = data_loader(
DATASET_PATH, 1, imgH, imgW, phase='test')
load_model(model_name, new_model)
test(new_model, device, test_loader, images, texts, lengths, converter, prediction_dir)
if __name__ == '__main__' :
main()
criterion = CTCLoss()
# cnn and rnn,define image, text, length placehoder
image = torch.FloatTensor(params.batchSize, 3, params.imgH, params.imgH)
text = torch.IntTensor(params.batchSize *
5) # text length no less than batchsize * 5
length = torch.IntTensor(params.batchSize) # define each word length
crnn = crnn.CRNN(params.imgH, nc, nclass, params.nh)
#if opt.cuda:
# crnn.cuda()
# image = image.cuda()
# criterion = criterion.cuda()
crnn = crnn.to(device)
image = image.to(device)
criterion = criterion.to(device)
text = text.to(device)
#length = length.to(device)
crnn.apply(weights_init) # self-define weight initialize function
#print("crnn =",crnn)
#print("crnn parameters =",crnn.cnn)
if params.crnn != '':
print('loading pretrained model from %s' % params.crnn)
crnn.load_state_dict(torch.load(params.crnn))
#for para in crnn.parameters():
#print("parameters =", para)#para.requires_grad=False
#crnn.state_dict()#get parameters list-------->cnn.conv0.weight cnn.conv0.bias cnn.conv1.weight cnn.conv1.bias cnn.conv2.weight cnn.conv2.bias cnn.batchnorm2.weight cnn.batchnorm2.bias
#crnn.state_dict().items()#get parameters name and its value ------------>('rnn.1.embedding.bias', tensor([-0.0602, -0.3962, -0.3687, -0.3052, -0.2965, -0.3442, -0.4302, -0.3631,\
#-0.3303, -0.2937, -0.2485, -0.4897, -0.2815, -0.3473, -0.3228, -0.2575,\
