update markdown docs (english) (#60)

colossalai/trainer/hooks/__init__.py

@@ -1,6 +1,7 @@
 from ._base_hook import BaseHook
 from ._checkpoint_hook import SaveCheckpointHook, LoadCheckpointHook
-from ._metric_hook import LossHook, Accuracy2DHook, AccuracyHook, MetricHook
+from ._metric_hook import (LossHook, Accuracy2DHook, AccuracyHook, MetricHook,
+                           Accuracy1DHook, Accuracy2p5DHook, Accuracy3DHook)
 from ._log_hook import LogMetricByEpochHook, TensorboardHook, LogTimingByEpochHook, LogMemoryByEpochHook
 from ._lr_scheduler_hook import LRSchedulerHook
 
@@ -8,6 +9,7 @@ __all__ = [
     'BaseHook', 'MetricHook',
     'LoadCheckpointHook', 'SaveCheckpointHook',
     'LossHook', 'AccuracyHook', 'Accuracy2DHook',
+    'Accuracy1DHook', 'Accuracy2p5DHook', 'Accuracy3DHook',
     'LogMetricByEpochHook', 'TensorboardHook', 'LogTimingByEpochHook', 'LogMemoryByEpochHook',
     'LRSchedulerHook'
 ]

colossalai/utils/data_sampler/data_parallel_sampler.py

@@ -108,7 +108,14 @@ class DataParallelSampler(Sampler):
         self.epoch = epoch
 
 
-def get_dataloader(dataset, shuffle=False, seed=1024, add_sampler=True, **kwargs):
+def get_dataloader(dataset,
+                   shuffle=False,
+                   seed=1024, 
+                   add_sampler=True, 
+                   drop_last=False,
+                   pin_memory=False,
+                   num_workers=0,
+                   **kwargs):
     '''Set up a deterministic dataloader (also configure seed workers, samplers and whether shuffle or not)
 
     .. note: when pipeline parallel is enabled, shuffle cannot be True
@@ -141,9 +148,15 @@ def get_dataloader(dataset, shuffle=False, seed=1024, add_sampler=True, **kwargs
         return DataLoader(dataset,
                           worker_init_fn=seed_worker,
                           shuffle=shuffle,
+                          drop_last=drop_last,
+                          pin_memory=pin_memory,
+                          num_workers=num_workers,
                           **_kwargs)
     else:
         return DataLoader(dataset,
                           sampler=sampler,
                           worker_init_fn=seed_worker,
+                          drop_last=drop_last,
+                          pin_memory=pin_memory,
+                          num_workers=num_workers,
                           **_kwargs)