Plug-ins overview
Unicore에는 총 5가지의 plug-in이 있다.
•
Model
Neural net 구조, parameter 정의
•
Loss (참고: fairseq 에서는 criterion)
loss function (output, target에 대해)
•
Task
주로 dictionary 형태로 object들을 저장하여
◦
Dataset loading, iterator helper function
◦
Model 초기화 helper function
◦
Criterion 초기화 helper function
◦
Loss 계산
•
Optimizer
Model parameter update
•
Learning rate scheduler
LR update
Training flow
•
unicore-train 이라는 command를 실행하면, 모든 plug-in들이 잘 정의되어 있을 때, high-level에서는 아래와 같은 방식으로 학습이 진행된다.
for epoch in range(num_epochs):
itr = task.get_batch_iterator(task.dataset('train'))
for num_updates, batch in enumerate(itr):
task.train_step(batch, model, loss, optimizer)
average_and_clip_gradients()
optimizer.step()
lr_scheduler.step_update(num_updates)
lr_scheduler.step(epoch)
Python
복사
◦
이때, task.train_step 의 high-level code는 아래와 같다.
def train_step(self, batch, model, loss, optimizer):
loss = criterion(model, batch)
optimizer.backward(loss)
return loss
Python
복사
Register new plug-ins
@register 라는 function decorator로 등록할 수 있다.
# unicore/model.py
@register_model('my_model')
class MyModel(BaseUnicoreModel):
(...)
#----------------------------------------#
# unicore/loss.py
@register_loss('my_loss')
class MyLoss(UnicoreLoss):
(...)
#----------------------------------------#
# unicore/task.py
@register_task('my_task')
class MyTask(UnicoreTask):
(...)
#----------------------------------------#
# unicore/optim/optim.py
@register_optimizer("my_opt")
class MyOpt(UnicoreOptimizer):
(...)
#----------------------------------------#
# unicore/optim/lr_scheduler/lr_scheduler.py
@register_lr_scheduler("my_lr_sche")
class MyLRScheduler(UnicoreLRScheduler):
(...)
#----------------------------------------#
Python
복사
→ 이렇게 등록한 plug-in 들은 Command-line Tool의 argument로 입력해서 선택, 사용할 수 있다.
예시)