add wandb

example/ae/standard/run.py

@@ -15,6 +15,8 @@ import deepke.attribution_extraction.standard.models as models
 from deepke.attribution_extraction.standard.tools import preprocess , CustomDataset, collate_fn ,train, validate
 from deepke.attribution_extraction.standard.utils import manual_seed, load_pkl
 
+import wandb
+
 
 logger = logging.getLogger(__name__)
 
@@ -26,6 +28,9 @@ def main(cfg):
     cfg.pos_size = 2 * cfg.pos_limit + 2
     logger.info(f'\n{cfg.pretty()}')
 
+    wandb.init(project="DeepKE_AE_Standard", name=cfg.model_name)
+    wandb.watch_called = False
+
     __Model__ = {
         'cnn': models.PCNN,
         'rnn': models.BiLSTM,
@@ -68,6 +73,8 @@ def main(cfg):
 
     model = __Model__[cfg.model_name](cfg)
     model.to(device)
+
+    wandb.watch(model, log="all")
     logger.info(f'\n {model}')
 
     optimizer = optim.Adam(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
@@ -95,6 +102,13 @@ def main(cfg):
 
         train_losses.append(train_loss)
         valid_losses.append(valid_loss)
+
+        wandb.log({
+            "train_loss":train_loss,
+            "valid_loss":valid_loss
+        })
+
+
         if best_f1 < valid_f1:
             best_f1 = valid_f1
             best_epoch = epoch
@@ -137,7 +151,12 @@ def main(cfg):
     logger.info('=====end of training====')
     logger.info('')
     logger.info('=====start test performance====')
-    validate(-1, model, test_dataloader, criterion, device, cfg)
+    _ , test_loss = validate(-1, model, test_dataloader, criterion, device, cfg)
+
+    wandb.log({
+        "test_loss":test_loss,
+    })
+    
     logger.info('=====ending====')
 
 

example/ner/few-shot/run.py

@@ -18,8 +18,10 @@ from deepke.name_entity_re.few_shot.module.mapping_type import mit_movie_mapping
 
 import warnings
 warnings.filterwarnings("ignore", category=UserWarning)
-from tensorboardX import SummaryWriter
-writer = SummaryWriter(log_dir='logs')
+
+import wandb
+
+writer = wandb.init(project="DeepKE_NER_Few")
 
 logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                     datefmt = '%m/%d/%Y %H:%M:%S',

example/ner/standard/run.py

@@ -19,6 +19,9 @@ import hydra
 from hydra import utils
 from deepke.name_entity_re.standard import *
 
+import wandb
+
+
 logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                     datefmt = '%m/%d/%Y %H:%M:%S',
                     level = logging.INFO)
@@ -53,7 +56,7 @@ class TrainNer(BertForTokenClassification):
         else:
             return logits
 
-
+wandb.init(project="DeepKE_NER_Standard")
 @hydra.main(config_path="conf", config_name='config')
 def main(cfg):
     
@@ -148,7 +151,9 @@ def main(cfg):
                     scheduler.step()  # Update learning rate schedule
                     model.zero_grad()
                     global_step += 1
-
+            wandb.log({
+                "train_loss":tr_loss/nb_tr_steps
+            })
         # Save a trained model and the associated configuration
         model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
         model_to_save.save_pretrained(utils.get_original_cwd()+'/'+cfg.output_dir)

example/re/document/run.py

@@ -12,6 +12,7 @@ from transformers.optimization import AdamW, get_linear_schedule_with_warmup
 
 from deepke.relation_extraction.document import *
 
+import wandb
 
 def train(args, model, train_features, dev_features, test_features):
     def logging(s, print_=True, log_=True):
@@ -87,6 +88,10 @@ def train(args, model, train_features, dev_features, test_features):
                         total_loss = 0
                         start_time = time.time()
 
+                        wandb.log({
+                            "train_loss":cur_loss
+                        })
+
                 if (step + 1) == len(train_dataloader) - 1 or (args.evaluation_steps > 0 and num_steps % args.evaluation_steps == 0 and step % args.gradient_accumulation_steps == 0):
                 # if step ==0:
                     logging('-' * 89)
@@ -96,11 +101,21 @@ def train(args, model, train_features, dev_features, test_features):
                     logging(
                         '| epoch {:3d} | time: {:5.2f}s | dev_result:{}'.format(epoch, time.time() - eval_start_time,
                                                                                 dev_output))
+
+                    wandb.log({
+                            "dev_result":dev_output
+                    })
+
                     logging('-' * 89)
                     if dev_score > best_score:
                         best_score = dev_score
                         logging(
                             '| epoch {:3d} | best_f1:{}'.format(epoch, best_score))
+
+                        wandb.log({
+                            "best_f1":best_score
+                        })
+
                         if args.save_path != "":
                             torch.save({
                                 'epoch': epoch,
@@ -170,10 +185,14 @@ def evaluate(args, model, features, tag="dev"):
         tag + "_re_r": re_r * 100,
         tag + "_average_loss": average_loss
     }
+
+    
+
     return best_f1, output
 
 
-
+wandb.init(project="DeepKE_RE_Document")
+wandb.watch_called = False
 @hydra.main(config_path="conf/config.yaml")
 def main(cfg):
     cwd = get_original_cwd()
@@ -207,6 +226,8 @@ def main(cfg):
         from_tf=bool(".ckpt" in cfg.model_name_or_path),
         config=config,
     )
+    wandb.watch(model, log="all")
+
 
     config.cls_token_id = tokenizer.cls_token_id
     config.sep_token_id = tokenizer.sep_token_id

example/re/few-shot/run.py

@@ -15,6 +15,8 @@ from tqdm import tqdm
 
 from deepke.relation_extraction.few_shot import *
 
+import wandb
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 
@@ -28,6 +30,8 @@ def logging(log_dir, s, print_=True, log_=True):
                 f_log.write(s + '\n')
 
 
+wandb.init(project="DeepKE_RE_Few")
+wandb.watch_called = False
 @hydra.main(config_path="conf/config.yaml")
 def main(cfg):
     cwd = get_original_cwd()
@@ -54,8 +58,11 @@ def main(cfg):
     #     model = torch.nn.DataParallel(model, device_ids = list(range(torch.cuda.device_count())))
 
     model.to(device)
+    wandb.watch(model, log="all")
 
     lit_model = BertLitModel(args=cfg, model=model, tokenizer=data.tokenizer)
+
+
     data.setup()
     
     if cfg.train_from_saved_model != '':
@@ -102,6 +109,11 @@ def main(cfg):
                     , print_=False)
                 log_loss = 0
         avrg_loss = total_loss / num_batch
+
+        wandb.log({
+                "train_loss": avrg_loss
+            })
+
         logging(cfg.log_dir,
             '| epoch {:2d} | train loss {:5.3f}'.format(
                 epoch, avrg_loss * 1000))
@@ -120,6 +132,12 @@ def main(cfg):
             logging(cfg.log_dir,
                 '| best_f1: {}'.format(best_f1))
             logging(cfg.log_dir,'-' * 89)
+
+            wandb.log({
+                "dev_result": f1,
+                "best_f1":best_f1
+            })
+            
             if cfg.save_path != "" and best != -1:
                 save_path = cfg.save_path
                 torch.save({

example/re/standard/run.py

@@ -15,6 +15,7 @@ import deepke.relation_extraction.standard.models as models
 from deepke.relation_extraction.standard.tools import preprocess , CustomDataset, collate_fn ,train, validate
 from deepke.relation_extraction.standard.utils import manual_seed, load_pkl
 
+import wandb
 
 logger = logging.getLogger(__name__)
 
@@ -26,6 +27,9 @@ def main(cfg):
     cfg.pos_size = 2 * cfg.pos_limit + 2
     logger.info(f'\n{cfg.pretty()}')
 
+    wandb.init(project="DeepKE_RE_Standard", name=cfg.model_name)
+    wandb.watch_called = False
+
     __Model__ = {
         'cnn': models.PCNN,
         'rnn': models.BiLSTM,
@@ -68,6 +72,8 @@ def main(cfg):
 
     model = __Model__[cfg.model_name](cfg)
     model.to(device)
+
+    wandb.watch(model, log="all")
     logger.info(f'\n {model}')
 
     optimizer = optim.Adam(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
@@ -95,6 +101,12 @@ def main(cfg):
 
         train_losses.append(train_loss)
         valid_losses.append(valid_loss)
+
+        wandb.log({
+            "train_loss":train_loss,
+            "valid_loss":valid_loss
+        })
+
         if best_f1 < valid_f1:
             best_f1 = valid_f1
             best_epoch = epoch
@@ -137,7 +149,12 @@ def main(cfg):
     logger.info('=====end of training====')
     logger.info('')
     logger.info('=====start test performance====')
-    validate(-1, model, test_dataloader, criterion, device, cfg)
+    _ , test_loss = validate(-1, model, test_dataloader, criterion, device, cfg)
+
+    wandb.log({
+        "test_loss":test_loss,
+    })
+    
     logger.info('=====ending====')
 
 

setup.py

@@ -1,7 +1,7 @@
 from setuptools import setup, find_packages
 setup(
     name='deepke',  # 打包后的包文件名
-    version='0.2.89',    #版本号
+    version='0.2.90',    #版本号
     keywords=["pip", "RE","NER","AE"],    # 关键字
     description='DeepKE 是基于 Pytorch 的深度学习中文关系抽取处理套件。',  # 说明
     long_description="client",  #详细说明
@@ -18,7 +18,6 @@ setup(
         'hydra-core==1.0.6',
         'tensorboard==2.4.1',
         'matplotlib==3.4.1',
-        'tensorboardX==2.4',
         'transformers==3.4.0',
         'jieba==0.42.1',
         'scikit-learn==0.24.1',
@@ -26,6 +25,7 @@ setup(
         'seqeval==1.2.2',
         'tqdm==4.60.0',
         'opt-einsum==3.3.0',
+        'wandb==0.12.7'
         "ujson"
     ], 
     classifiers=[

src/deepke/name_entity_re/few_shot/module/train.py

@@ -67,9 +67,11 @@ class Trainer(object):
                         avg_loss = float(avg_loss) / self.refresh_step
                         print_output = "loss:{:<6.5f}".format(avg_loss)
                         pbar.update(1)
-                        pbar.set_postfix_str(print_output)
-                        self.writer.add_scalar(tag='loss', scalar_value=avg_loss, global_step=self.step)    # tensorbordx
+                        pbar.set_postfix_str(print_output)            
+                        # self.writer.add_scalar(tag='loss', scalar_value=avg_loss, global_step=self.step)    # tensorbordx
+                        self.writer.log({ 'avg_loss':avg_loss})
                         avg_loss = 0
+                        
                 if epoch >= self.eval_begin_epoch:
                     self.evaluate(epoch)   # generator to dev.
             pbar.close()
@@ -95,9 +97,11 @@ class Trainer(object):
                 pbar.close()
                 self.logger.info("Epoch {}/{}, best f1: {}, current f1 score: {:.2f}, recall: {:.2f}, precision: {:.2f}."\
                             .format(epoch, self.num_epochs, self.best_metric, eva_result['f'], eva_result['rec'], eva_result['pre']))
-                self.writer.add_scalars('evaluate', {'f1': eva_result['f'],
-                                                    'recall': eva_result['rec'],
-                                                    'precision': eva_result['pre']}, epoch)
+                # self.writer.add_scalars('evaluate', {'f1': eva_result['f'],
+                #                                     'recall': eva_result['rec'],
+                #                                     'precision': eva_result['pre']}, epoch)
+                self.writer.log({'f1': eva_result['f'],'recall': eva_result['rec'],'precision': eva_result['pre']})
+
                 if eva_result['f'] >= self.best_metric:  # this epoch get best performance
                     self.logger.info("Get better performance at epoch {}".format(epoch))
                     self.best_dev_epoch = epoch