Merge pull request #14 from codertimo/alpha0.0.1a4

alpha-0.0.1a4 version released

LICENSE

@@ -186,7 +186,7 @@
  same "printed page" as the copyright notice for easier
  identification within third-party archives.
 
- Copyright 2018 Junseong Kim, Scatter Labs, BERT contributors
+ Copyright 2018 Junseong Kim, Scatter Lab, BERT contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

README.md

@@ -60,7 +60,7 @@ bert-vocab -c data/corpus.small -o data/vocab.small
 
 ### 2. Train your own BERT model
 ```shell
-bert -c data/dataset.small -v data/vocab.small -o output/bert.model
+bert -c data/corpus.small -v data/vocab.small -o output/bert.model
 ```
 
 ## Language Model Pre-training

bert_pytorch/__main__.py

@@ -10,27 +10,30 @@
 def train():
  parser = argparse.ArgumentParser()
 
- parser.add_argument("-c", "--train_dataset", required=True, type=str)
- parser.add_argument("-t", "--test_dataset", type=str, default=None)
- parser.add_argument("-v", "--vocab_path", required=True, type=str)
- parser.add_argument("-o", "--output_path", required=True, type=str)
-
- parser.add_argument("-hs", "--hidden", type=int, default=256)
- parser.add_argument("-l", "--layers", type=int, default=8)
- parser.add_argument("-a", "--attn_heads", type=int, default=8)
- parser.add_argument("-s", "--seq_len", type=int, default=20)
-
- parser.add_argument("-b", "--batch_size", type=int, default=64)
- parser.add_argument("-e", "--epochs", type=int, default=10)
- parser.add_argument("-w", "--num_workers", type=int, default=5)
- parser.add_argument("--with_cuda", type=bool, default=True)
- parser.add_argument("--log_freq", type=int, default=10)
- parser.add_argument("--corpus_lines", type=int, default=None)
-
- parser.add_argument("--lr", type=float, default=1e-3)
- parser.add_argument("--adam_weight_decay", type=float, default=0.01)
- parser.add_argument("--adam_beta1", type=float, default=0.9)
- parser.add_argument("--adam_beta2", type=float, default=0.999)
+ parser.add_argument("-c", "--train_dataset", required=True, type=str, help="train dataset for train bert")
+ parser.add_argument("-t", "--test_dataset", type=str, default=None, help="test set for evaluate train set")
+ parser.add_argument("-v", "--vocab_path", required=True, type=str, help="built vocab model path with bert-vocab")
+ parser.add_argument("-o", "--output_path", required=True, type=str, help="ex)output/bert.model")
+
+ parser.add_argument("-hs", "--hidden", type=int, default=256, help="hidden size of transformer model")
+ parser.add_argument("-l", "--layers", type=int, default=8, help="number of layers")
+ parser.add_argument("-a", "--attn_heads", type=int, default=8, help="number of attention heads")
+ parser.add_argument("-s", "--seq_len", type=int, default=20, help="maximum sequence len")
+
+ parser.add_argument("-b", "--batch_size", type=int, default=64, help="number of batch_size")
+ parser.add_argument("-e", "--epochs", type=int, default=10, help="number of epochs")
+ parser.add_argument("-w", "--num_workers", type=int, default=5, help="dataloader worker size")
+
+ parser.add_argument("--with_cuda", type=bool, default=True, help="training with CUDA: true, or false")
+ parser.add_argument("--log_freq", type=int, default=10, help="printing loss every n iter: setting n")
+ parser.add_argument("--corpus_lines", type=int, default=None, help="total number of lines in corpus")
+ parser.add_argument("--cuda_devices", type=int, nargs='+', default=None, help="CUDA device ids")
+ parser.add_argument("--on_memory", type=bool, default=True, help="Loading on memory: true or false")
+
+ parser.add_argument("--lr", type=float, default=1e-3, help="learning rate of adam")
+ parser.add_argument("--adam_weight_decay", type=float, default=0.01, help="weight_decay of adam")
+ parser.add_argument("--adam_beta1", type=float, default=0.9, help="adam first beta value")
+ parser.add_argument("--adam_beta2", type=float, default=0.999, help="adam first beta value")
 
  args = parser.parse_args()
 
@@ -39,11 +42,12 @@ def train():
  print("Vocab Size: ", len(vocab))
 
  print("Loading Train Dataset", args.train_dataset)
- train_dataset = BERTDataset(args.train_dataset, vocab, seq_len=args.seq_len, corpus_lines=args.corpus_lines)
+ train_dataset = BERTDataset(args.train_dataset, vocab, seq_len=args.seq_len,
+ corpus_lines=args.corpus_lines, on_memory=args.on_memory)
 
  print("Loading Test Dataset", args.test_dataset)
- test_dataset = BERTDataset(args.test_dataset, vocab,
-  seq_len=args.seq_len) if args.test_dataset is not None else None
+ test_dataset = BERTDataset(args.test_dataset, vocab, seq_len=args.seq_len, on_memory=args.on_memory) \
+ if args.test_dataset is not None else None
 
  print("Creating Dataloader")
  train_data_loader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.num_workers)
@@ -56,7 +60,7 @@ def train():
  print("Creating BERT Trainer")
  trainer = BERTTrainer(bert, len(vocab), train_dataloader=train_data_loader, test_dataloader=test_data_loader,
  lr=args.lr, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay,
- with_cuda=args.with_cuda, log_freq=args.log_freq)
+ with_cuda=args.with_cuda, cuda_devices=args.cuda_devices, log_freq=args.log_freq)
 
  print("Training Start")
  for epoch in range(args.epochs):

bert_pytorch/dataset/dataset.py

@@ -5,19 +5,37 @@
 
 
 class BERTDataset(Dataset):
- def __init__(self, corpus_path, vocab, seq_len, encoding="utf-8", corpus_lines=None):
+ def __init__(self, corpus_path, vocab, seq_len, encoding="utf-8", corpus_lines=None, on_memory=True):
  self.vocab = vocab
  self.seq_len = seq_len
 
+ self.on_memory = on_memory
+ self.corpus_lines = corpus_lines
+ self.corpus_path = corpus_path
+ self.encoding = encoding
+
  with open(corpus_path, "r", encoding=encoding) as f:
- self.datas = [line[:-1].split("\t")
- for line in tqdm.tqdm(f, desc="Loading Dataset", total=corpus_lines)]
+ if self.corpus_lines is None and not on_memory:
+ for _ in tqdm.tqdm(f, desc="Loading Dataset", total=corpus_lines):
+ self.corpus_lines += 1
+
+ if on_memory:
+ self.lines = [line[:-1].split("\t")
+ for line in tqdm.tqdm(f, desc="Loading Dataset", total=corpus_lines)]
+ self.corpus_lines = len(self.lines)
+
+ if not on_memory:
+ self.file = open(corpus_path, "r", encoding=encoding)
+ self.random_file = open(corpus_path, "r", encoding=encoding)
+
+ for _ in range(random.randint(self.corpus_lines if self.corpus_lines < 1000 else 1000)):
+ self.random_file.__next__()
 
  def __len__(self):
- return len(self.datas)
+ return self.corpus_lines
 
  def __getitem__(self, item):
- t1, (t2, is_next_label) = self.datas[item][0], self.random_sent(item)
+ t1, t2, is_next_label = self.random_sent(item)
  t1_random, t1_label = self.random_word(t1)
  t2_random, t2_label = self.random_word(t2)
 
@@ -49,16 +67,18 @@ def random_word(self, sentence):
  for i, token in enumerate(tokens):
  prob = random.random()
  if prob < 0.15:
- # 80% randomly change token to make token
- if prob < prob * 0.8:
+ prob /= 0.15
+
+ # 80% randomly change token to mask token
+ if prob < 0.8:
  tokens[i] = self.vocab.mask_index
 
  # 10% randomly change token to random token
- elif prob * 0.8 <= prob < prob * 0.9:
+ elif prob < 0.9:
  tokens[i] = random.randrange(len(self.vocab))
 
  # 10% randomly change token to current token
- elif prob >= prob * 0.9:
+ else:
  tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index)
 
  output_label.append(self.vocab.stoi.get(token, self.vocab.unk_index))
@@ -70,8 +90,36 @@ def random_word(self, sentence):
  return tokens, output_label
 
  def random_sent(self, index):
+ t1, t2 = self.get_corpus_line(index)
+
  # output_text, label(isNotNext:0, isNext:1)
  if random.random() > 0.5:
- return self.datas[index][1], 1
+ return t1, t2, 1
+ else:
+ return t1, self.get_random_line(), 0
+
+ def get_corpus_line(self, item):
+ if self.on_memory:
+ return self.lines[item][0], self.lines[item][1]
  else:
- return self.datas[random.randrange(len(self.datas))][1], 0
+ line = self.file.__next__()
+ if line is None:
+ self.file.close()
+ self.file = open(self.corpus_path, "r", encoding=self.encoding)
+ line = self.file.__next__()
+
+ t1, t2 = line[:-1].split("\t")
+ return t1, t2
+
+ def get_random_line(self):
+ if self.on_memory:
+ return self.lines[random.randrange(len(self.lines))][1]
+
+ line = self.file.__next__()
+ if line is None:
+ self.file.close()
+ self.file = open(self.corpus_path, "r", encoding=self.encoding)
+ for _ in range(random.randint(self.corpus_lines if self.corpus_lines < 1000 else 1000)):
+ self.random_file.__next__()
+ line = self.random_file.__next__()
+ return line[:-1].split("\t")[1]

bert_pytorch/model/embedding/position.py

@@ -13,7 +13,7 @@ def __init__(self, d_model, max_len=512):
  pe.require_grad = False
 
  position = torch.arange(0, max_len).float().unsqueeze(1)
- div_term = (torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)).float().exp()
+ div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
 
  pe[:, 0::2] = torch.sin(position * div_term)
  pe[:, 1::2] = torch.cos(position * div_term)

bert_pytorch/trainer/optim_schedule.py

@@ -0,0 +1,35 @@
+'''A wrapper class for optimizer '''
+import numpy as np
+
+
+class ScheduledOptim():
+ '''A simple wrapper class for learning rate scheduling'''
+
+ def __init__(self, optimizer, d_model, n_warmup_steps):
+ self._optimizer = optimizer
+ self.n_warmup_steps = n_warmup_steps
+ self.n_current_steps = 0
+ self.init_lr = np.power(d_model, -0.5)
+
+ def step_and_update_lr(self):
+ "Step with the inner optimizer"
+ self._update_learning_rate()
+ self._optimizer.step()
+
+ def zero_grad(self):
+ "Zero out the gradients by the inner optimizer"
+ self._optimizer.zero_grad()
+
+ def _get_lr_scale(self):
+ return np.min([
+ np.power(self.n_current_steps, -0.5),
+ np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])
+
+ def _update_learning_rate(self):
+ ''' Learning rate scheduling per step '''
+
+ self.n_current_steps += 1
+ lr = self.init_lr * self._get_lr_scale()
+
+ for param_group in self._optimizer.param_groups:
+ param_group['lr'] = lr

bert_pytorch/trainer/pretrain.py

@@ -4,6 +4,7 @@
 from torch.utils.data import DataLoader
 
 from ..model import BERTLM, BERT
+from .optim_schedule import ScheduledOptim
 
 import tqdm
 
@@ -21,8 +22,8 @@ class BERTTrainer:
 
  def __init__(self, bert: BERT, vocab_size: int,
  train_dataloader: DataLoader, test_dataloader: DataLoader = None,
- lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01,
- with_cuda: bool = True, log_freq: int = 10):
+ lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01, warmup_steps=10000,
+ with_cuda: bool = True, cuda_devices=None, log_freq: int = 10):
  """
  :param bert: BERT model which you want to train
  :param vocab_size: total word vocab size
@@ -45,16 +46,17 @@ def __init__(self, bert: BERT, vocab_size: int,
  self.model = BERTLM(bert, vocab_size).to(self.device)
 
  # Distributed GPU training if CUDA can detect more than 1 GPU
- if torch.cuda.device_count() > 1:
+ if with_cuda and torch.cuda.device_count() > 1:
  print("Using %d GPUS for BERT" % torch.cuda.device_count())
- self.model = nn.DataParallel(self.model)
+ self.model = nn.DataParallel(self.model, device_ids=cuda_devices)
 
  # Setting the train and test data loader
  self.train_data = train_dataloader
  self.test_data = test_dataloader
 
  # Setting the Adam optimizer with hyper-param
  self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay)
+ self.optim_schedule = ScheduledOptim(self.optim, self.bert.hidden, n_warmup_steps=warmup_steps)
 
  # Using Negative Log Likelihood Loss function for predicting the masked_token
  self.criterion = nn.NLLLoss(ignore_index=0)
@@ -110,9 +112,9 @@ def iteration(self, epoch, data_loader, train=True):
 
  # 3. backward and optimization only in train
  if train:
- self.optim.zero_grad()
+ self.optim_schedule.zero_grad()
  loss.backward()
- self.optim.step()
+ self.optim_schedule.step_and_update_lr()
 
  # next sentence prediction accuracy
  correct = next_sent_output.argmax(dim=-1).eq(data["is_next"]).sum().item()

requirements.txt

@@ -1,3 +1,3 @@
 tqdm
 numpy
-torch>=0.4.0
+torch>=0.4.0

setup.py

@@ -3,10 +3,10 @@
 import os
 import sys
 
-__version__ = "0.0.1a3"
+__version__ = "0.0.1a4"
 
 with open("requirements.txt") as f:
- require_packages = [line[:-1] for line in f]
+ require_packages = [line[:-1] if line[-1] == "\n" else line for line in f]
 
 with open("README.md", "r", encoding="utf-8") as f:
  long_description = f.read()