tensorflow에서는 CTC loss가 내장되어 있어 CTC를 구현할 필요가 없었다.
pytorch에는 없는데 구현하기에는 까다롭고 필요성을 못느껴
Listen-Attend-and-Spell이라는 아키텍처의 OCR을 만들고자 한다.
물론, 이 글에서도 Recognition model만 해보고자 한다.
Listen-Attend-and-Spell
- Listen-Attend-and-Spell(LAS) 모델은 음성 인식 분야에서 사용되는 순환 신경망 기반 모델이다.
- End-to-End 방식으로 학습이 가능하여 사전 정의된 작업 단위 없이 음성을 직접 문자열로 변환할 수 있다.
- 이 모델은 세 가지 주요 단계로 구성되어 있다.
- Listen: 음성 입력을 받아 특징을 추출하는 단계입니다. 일반적으로 CNN, RNN 등의 모델이 사용된다.
- Attend: 추출된 특징에 주의(Attention) 메커니즘을 적용하여 관련 부분에 집중하는 단계이다.
- Spell: 주의 집중된 특징을 통해 문자 시퀀스를 생성하는 단계이다. 디코더로 작동한다.
모델 구축
- LAS 모델 처럼 구축할 OCR 모델을 Encoder-Decoder 로 나눈다.
- LAS 모델과 다르게 입력은 이미지로 들어가므로 Encoder 모델은 CNN과 RNN을 조합한 형태로 만든다.
- Decoder 모델은 RNN과 Attention 메커니즘을 조합한 형태로 만들었다.
인코더
- RNN을 추가하고 양방향으로 바꾸고 나서 학습이 잘 진행 되었었다.
- 코드에서 RNN의 파라미터에서 batch_first=True
- CNN out의 채널 위치와 RNN out의 채널 위치가 다르기에 하단에 transpose를 주었다.
class ConvEncoder(nn.Module):
def __init__(self, input_dim, output_dim=128):
super().__init__()
self.conv2d_1 = nn.Conv2d(input_dim, 32, kernel_size=3, padding=1)
self.conv2d_2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.conv2d_3 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
self.conv1d_1 = nn.Conv1d(1024, 256, kernel_size=3, padding=1)
self.conv1d_2 = nn.Conv1d(256, 128, kernel_size=1)
self.lstm = nn.LSTM(128, output_dim//2, batch_first=True, bidirectional=True)
self.batch_norm2d = nn.BatchNorm2d(64)
self.batch_norm1d = nn.BatchNorm1d(output_dim)
self.max_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.activation = nn.LeakyReLU()
self.dropout2d = nn.Dropout2d(p=0.1) # 채널에 대하여 dropout
self.dropout1d = nn.Dropout1d(p=0.1) # 채널에 대하여 dropout
nn.init.kaiming_normal_(self.conv2d_1.weight)
nn.init.kaiming_normal_(self.conv2d_2.weight)
nn.init.kaiming_normal_(self.conv2d_3.weight)
nn.init.kaiming_normal_(self.conv1d_1.weight)
nn.init.kaiming_normal_(self.conv1d_2.weight)
nn.init.xavier_uniform_(self.lstm.weight_ih_l0)
nn.init.xavier_uniform_(self.lstm.weight_hh_l0)
def forward(self, images):
# (batch, 3, 64, 512)
x = self.conv2d_1(images)
x = self.activation(x)
x = self.max_pool(x)
# (batch, 32, 32, 256)
x = self.conv2d_2(x)
x = self.activation(x)
x = self.max_pool(x)
# (batch, 64, 16, 128)
x = self.conv2d_3(x)
x = self.batch_norm2d(x)
x = self.activation(x)
x = self.dropout2d(x)
# (batch, 64, 16, 128)
x = x.view(x.size(0), -1, x.size(-1))
# (batch, 1024, 128)
x = self.conv1d_1(x)
x = self.conv1d_2(x)
x = self.batch_norm1d(x)
x = self.activation(x)
x = self.dropout1d(x)
# (batch, 128, 128)
x = x.transpose(2,1)
# (batch, 128, 128)
x, _ = self.lstm(x)
# (batch, 128, 64*2)
return x디코더
- 이 코드에서 beam_search 메서드 부분은 페이지 절약을 위해 여기서는 지웠다.
- 자세한 코드는 글 아래 링크 참조
- loss로 NLL을 사용하기에 out을 logsoftamx를 사용하였다.
- loss로 CrossEntropy를 사용한다면 그냥 linear로 out을 사용해도 될 것이다.
- logsoftamx를 사용하면 beam_search 에서 score 계산하기 편하기도 하다.
- ground_truth가 주어지면 forward, 주어지지 않으면 기본으로 greedy_search로 진행한다.
import numpy as np
import torch.nn.functional as F
class Speller(nn.Module):
def __init__(self, encoder_dim, target_dim, hidden_dim=128,
sos_id=13, eos_id=14, pad_id=15, max_len=20):
super().__init__()
self.rnn_layer = nn.LSTM(hidden_dim*2, hidden_dim, batch_first=True, num_layers=2, dropout=0.5)
self.attention = nn.MultiheadAttention(hidden_dim, 4, batch_first=True)
self.convertor_linear = nn.Linear(encoder_dim, hidden_dim)
self.output_linear = nn.Linear(hidden_dim*2, target_dim)
self.logsoftmax = nn.LogSoftmax(dim=-1)
self.dropout = nn.Dropout(0.2)
self.emb = nn.Embedding(16, hidden_dim)
self.target_dim = target_dim
self.hidden_dim = hidden_dim
self.sos_id = sos_id
self.eos_id = eos_id
self.pad_id = pad_id
self.max_len = max_len
nn.init.xavier_uniform_(self.rnn_layer.weight_ih_l0)
nn.init.xavier_uniform_(self.rnn_layer.weight_hh_l0)
nn.init.xavier_uniform_(self.rnn_layer.weight_ih_l1)
nn.init.xavier_uniform_(self.rnn_layer.weight_hh_l1)
nn.init.xavier_uniform_(self.convertor_linear.weight)
nn.init.xavier_uniform_(self.output_linear.weight)
nn.init.xavier_uniform_(self.attention.in_proj_weight)
nn.init.xavier_uniform_(self.attention.out_proj.weight)
def forward_step(self, rnn_in, converted, hidden_state):
# rnn_input : (batch, 1, hidden_dim*2)
# converted : (batch, encoder_seq, hidden_dim)
# rnn_out : (batch, 1, hidden_dim)
# context : (batch, 1, hidden_dim)
# att_score : (batch, 1, encoder_seq)
# concat_out : (batch, 1, hidden_dim*2)
# step_out : (batch, 1, target_dim)
rnn_out, hidden_state = self.rnn_layer(rnn_in, hidden_state)
context, att_score = self.attention(rnn_out, converted, converted)
concat_out = torch.cat([rnn_out, context], dim=-1)
x = self.dropout(concat_out)
x = self.output_linear(x)
step_out = self.logsoftmax(x)
return step_out, hidden_state, context, att_score
def greedy_search(self, encoder_outputs):
device = encoder_outputs.device
batch_size = encoder_outputs.size(0)
converted = self.convertor_linear(encoder_outputs).tanh() # (batch, encoder_seq, hidden_dim)
sos = torch.tensor(self.sos_id).tile(batch_size,1).to(device)
rnn_in = torch.cat([self.emb(sos).tanh(), converted[:,0:1,:]], dim=-1) # (batch, 1, dim*2)
hidden_state = None
prob_log_seq = []
target_idxs = []
att_recode = []
for i in range(self.max_len):
step_out, hidden_state, context, att_score = self.forward_step(rnn_in, converted, hidden_state)
argmax = step_out.max(dim=-1)[1] # (batch, 1)
rnn_in = torch.cat([self.emb(argmax).tanh(), context], dim=-1) # (batch, 1, hidden_dim*2)
prob_log_seq.append(step_out) # (batch, 1, target_dim)
target_idxs.append(argmax)
att_recode.append(att_score) # (batch, 1, encoder_seq)
prob_log_seq = torch.cat(prob_log_seq, dim=1) # (batch, max_len, dim)
target_idxs = torch.cat(target_idxs, dim=1) # (batch, max_len)
att_recode = torch.cat(att_recode, dim=1) # (batch, max_len, encoder_seq)
return prob_log_seq, target_idxs, att_recode
def forward(self, encoder_outputs, target_idxs=None):
if target_idxs is None:
return self.greedy_search(encoder_outputs)
# encoder_outputs : (batch, encoder_seq, encoder_dim)
# target_idxs : (batch, max_len)
device = encoder_outputs.device
batch_size = encoder_outputs.size(0)
converted = self.convertor_linear(encoder_outputs).tanh() # (batch, encoder_seq, hidden_dim)
sos = torch.tensor(self.sos_id).tile(batch_size,1).to(device)
rnn_in = torch.cat([self.emb(sos).tanh(), converted[:,0:1,:]], dim=-1) # (batch, 1, dim*2)
hidden_state = None
prob_log_seq = []
att_recode = []
for i in range(self.max_len):
step_out, hidden_state, context, att_score = self.forward_step(rnn_in, converted, hidden_state)
rnn_in = torch.cat([self.emb(target_idxs[:, i:i+1]).tanh(), context], dim=-1) # (batch, 1, hidden_dim*2)
prob_log_seq.append(step_out) # (batch, 1, target_dim)
att_recode.append(att_score) # (batch, 1, encoder_seq)
prob_log_seq = torch.cat(prob_log_seq, dim=1) # (batch, max_len, dim)
att_recode = torch.cat(att_recode, dim=1) # (batch, max_len, encoder_seq)
return prob_log_seq, target_idxs, att_recode모델 학습
- loss는 NLL(Negative-Log-Likelihood) loss를 사용하였고, class weight를 주었다.
- optimizer로 AdamW를 사용.
- scheduler로 ReduceLROnPlateau.
- early stopping을 넣었다.
- LAS의 특징으로 0.1 확률로 모델에 ground_truth가 주어지지 않는다.
- 추론환경과 학습환경 간의 불일치가 일어나서 성능저하 예방을 위해 필요하다고 한다.
from torch import optim
import numpy as np
# torch.autograd.set_detect_anomaly(True)
EPOCHS = 300
# class weight
_, cnts = np.unique(y_train.reshape(-1).tolist()+[SOS_ID], return_counts=True)
weight = torch.from_numpy(np.log(np.sum(cnts) / cnts)).float().cuda()
# 모델, 손실 함수, 옵티마이저 초기화
model = ImageLAS(sos_id=SOS_ID, eos_id=EOS_ID, pad_id=PAD_ID, max_len=MAX_LENGTH).cuda()
loss_fn = nn.NLLLoss(weight=weight, ignore_index=SOS_ID).cuda() # CrossEntropyLoss().cuda() # (logit, target)
optimizer = optim.AdamW(model.parameters(), lr=0.002)
# scheduler = CosineDecayRestarts(optimizer)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.5, min_lr=1e-5, patience=20,)
# Early stopping을 위한 변수 초기화
best_val_loss = float('inf')
patience = 50 # 일정 기간 동안 성능이 개선되지 않으면 학습 중지
early_stopping_counter = 0
# history
loss_history = []
val_loss_history = []
# 학습 루틴
for epoch in range(EPOCHS):
# 학습
model.train() # 모델을 학습 모드로 설정
tra_loss = 0.0
for inputs, target_idxs in train_dataloader:
# for inputs, target_idxs in tqdm(train_dataloader, total=len(train_dataloader)):
inputs, target_idxs = inputs.cuda(), target_idxs.cuda()
optimizer.zero_grad() # 그래디언트 초기화
if 0.9 < np.random.random(): target_idxs = None
outputs, target_idxs, _ = model(inputs, target_idxs) # 모델 출력 계산
loss = loss_fn(outputs.permute(0,2,1), target_idxs) # 손실 계산
loss.backward() # 역전파 수행
optimizer.step() # 최적화 수행
tra_loss += loss.item()
# 검증
model.eval() # 모델을 평가 모드로 설정
val_loss = 0.0
with torch.no_grad():
for inputs, target_idxs in valid_dataloader:
inputs, target_idxs = inputs.cuda(), target_idxs.cuda()
outputs, _, _ = model(inputs, target_idxs)
val_loss += loss_fn(outputs.permute(0,2,1), target_idxs).item() # 검증 손실 누적
# 출력
tra_loss /= len(train_dataloader)
val_loss /= len(valid_dataloader)
loss_history.append(tra_loss)
val_loss_history.append(val_loss)
if (epoch+1)%5 == 0:
print(f"Epoch {epoch+1:03d}\tTrain Loss: {tra_loss:.6f}" +
f"\tValid Loss: {val_loss:.6f}\tlr: {optimizer.param_groups[0]['lr']:.6f}")
# 스케줄러에게 현재 검증 손실을 전달하여 학습률 갱신
scheduler.step(val_loss)
# 검증 손실이 이전보다 크면 early stopping counter를 증가시킴
if val_loss > best_val_loss:
early_stopping_counter += 1
else:
best_val_loss = val_loss
early_stopping_counter = 0
if best_val_loss < 0.35:
torch.save(model.state_dict(), './models/best.pt')
# early stopping 조건 충족 시 학습 중지
if early_stopping_counter >= patience:
print("Early stopping! No improvement in validation loss.")
print(f"Epoch {epoch+1:03d}\tTrain Loss: {tra_loss:.6f}" +
f"\tValid Loss: {val_loss:.6f}\tlr: {optimizer.param_groups[0]['lr']:.6f}")
print(f"Best_val_loss : {best_val_loss:.6f}")
break
model = ImageLAS(sos_id=SOS_ID, eos_id=EOS_ID, pad_id=PAD_ID, max_len=MAX_LENGTH)
model.load_state_dict(torch.load('./models/best.pt'))
print("Restored to best model.")History 그래프
- Best loss :
Train Loss: 0.003741
Valid Loss: 0.038099
모델 테스트
Attention Score 시각화
