Add grouped cache for data bigger than CPU RAM

map2map/args.py

@@ -1,4 +1,5 @@
 import argparse
+import warnings
 
 from .train import ckpt_link
 
@@ -53,7 +54,7 @@ def add_common_args(parser):
     parser.add_argument('--load-state', default=ckpt_link, type=str,
             help='path to load the states of model, optimizer, rng, etc. '
             'Default is the checkpoint. '
-            'Start from scratch if the checkpoint does not exist')
+            'Start from scratch if set empty or the checkpoint is missing')
     parser.add_argument('--load-state-non-strict', action='store_false',
             help='allow incompatible keys when loading model states',
             dest='load_state_strict')
@@ -173,6 +174,17 @@ def set_common_args(args):
         if args.batches > 1:
             args.loader_workers = args.batches
 
+    if not args.cache and args.cache_maxsize is not None:
+        args.cache_maxsize = None
+        warnings.warn('Resetting cache maxsize given cache is disabled',
+                      RuntimeWarning)
+    if (args.cache and args.cache_maxsize is not None
+            and args.cache_maxsize < 1):
+        args.cache = False
+        args.cache_maxsize = None
+        warnings.warn('Disabling cache given cache maxsize < 1',
+                      RuntimeWarning)
+
 
 def set_train_args(args):
     set_common_args(args)

map2map/data/__init__.py

@@ -1 +1,2 @@
 from .fields import FieldDataset
+from .sampler import GroupedRandomSampler

map2map/data/fields.py

@@ -11,7 +11,7 @@ from .norms import import_norm
 class FieldDataset(Dataset):
     """Dataset of lists of fields.
 
-    `in_patterns` is a list of glob patterns for the input fields.
+    `in_patterns` is a list of glob patterns for the input field files.
     For example, `in_patterns=['/train/field1_*.npy', '/train/field2_*.npy']`.
     Likewise `tgt_patterns` is for target fields.
     Input and target fields are matched by sorting the globbed files.
@@ -51,15 +51,9 @@ class FieldDataset(Dataset):
 
         assert len(self.in_files) == len(self.tgt_files), \
                 'number of input and target fields do not match'
+        self.nfile = len(self.in_files)
 
-        assert len(self.in_files) > 0, 'file not found'
+        assert self.nfile > 0, 'file not found'
-
-        if div_data:
-            files = len(self.in_files) // world_size
-            self.in_files = self.in_files[rank * files : (rank + 1) * files]
-            self.tgt_files = self.tgt_files[rank * files : (rank + 1) * files]
-
-        assert len(self.in_files) > 0, 'files not divisible among ranks'
 
         self.in_chan = [np.load(f).shape[0] for f in self.in_files[0]]
         self.tgt_chan = [np.load(f).shape[0] for f in self.tgt_files[0]]
@@ -92,7 +86,7 @@ class FieldDataset(Dataset):
         else:
             self.crop = np.broadcast_to(crop, self.size.shape)
             self.reps = self.size // self.crop
-        self.tot_reps = int(np.prod(self.reps))
+        self.ncrop = int(np.prod(self.reps))
 
         assert isinstance(pad, int), 'only support symmetric padding for now'
         self.pad = np.broadcast_to(pad, (self.ndim, 2))
@@ -104,19 +98,44 @@ class FieldDataset(Dataset):
         if cache:
             self.get_fields = lru_cache(maxsize=cache_maxsize)(self.get_fields)
 
+        if div_data:
+            self.samples = []
+
+            # first add full fields when num_fields > num_GPU
+            for i in range(rank, self.nfile // world_size * world_size,
+                           world_size):
+                self.samples.append(
+                    range(i * self.ncrop, (i + 1) * self.ncrop)
+                )
+
+            # then split the rest into fractions of fields
+            # drop the last incomplete batch of samples
+            frac_start = self.nfile // world_size * world_size * self.ncrop
+            frac_samples = self.nfile % world_size * self.ncrop // world_size
+            self.samples.append(range(frac_start + rank * frac_samples,
+                                      frac_start + (rank + 1) * frac_samples))
+
+            self.samples = np.concatenate(self.samples)
+        else:
+            self.samples = np.arange(self.nfile * self.ncrop)
+        self.nsample = len(self.samples)
+
+        self.rank = rank
+
     def get_fields(self, idx):
         in_fields = [np.load(f) for f in self.in_files[idx]]
         tgt_fields = [np.load(f) for f in self.tgt_files[idx]]
         return in_fields, tgt_fields
 
     def __len__(self):
-        return len(self.in_files) * self.tot_reps
+        return self.nsample
 
     def __getitem__(self, idx):
-        idx, sub_idx = idx // self.tot_reps, idx % self.tot_reps
+        idx = self.samples[idx]
-        start = np.unravel_index(sub_idx, self.reps) * self.crop
 
-        in_fields, tgt_fields = self.get_fields(idx)
+        in_fields, tgt_fields = self.get_fields(idx // self.ncrop)
+
+        start = np.unravel_index(idx % self.ncrop, self.reps) * self.crop
 
         in_fields = crop(in_fields, start, self.crop, self.pad)
         tgt_fields = crop(tgt_fields, start * self.scale_factor,

map2map/data/sampler.py

@@ -0,0 +1,31 @@
+from itertools import chain
+import torch
+from torch.utils.data import Sampler
+
+
+class GroupedRandomSampler(Sampler):
+    """Sample randomly within each group of samples and sequentially from group
+    to group.
+
+    This behaves like a simple random sampler by default
+    """
+    def __init__(self, data_source, group_size=None):
+        self.data_source = data_source
+        self.sample_size = len(data_source)
+
+        if group_size is None:
+            group_size = self.sample_size
+        self.group_size = group_size
+
+    def __iter__(self):
+        starts = range(0, self.sample_size, self.group_size)
+        sizes = [self.group_size] * (len(starts) - 1)
+        sizes.append(self.sample_size - starts[-1])
+
+        return iter(chain(*[
+            (start + torch.randperm(size)).tolist()
+            for start, size in zip(starts, sizes)
+        ]))
+
+    def __len__(self):
+        return self.sample_size

map2map/train.py

@@ -13,7 +13,7 @@ from torch.utils.data.distributed import DistributedSampler
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
 
-from .data import FieldDataset
+from .data import FieldDataset, GroupedRandomSampler
 from .data.figures import fig3d
 from . import models
 from .models import (narrow_like,
@@ -48,11 +48,13 @@ def gpu_worker(local_rank, node, args):
     device = torch.device('cuda', local_rank)
     torch.cuda.device(device)
 
-    torch.manual_seed(args.seed)
-    #torch.backends.cudnn.deterministic = True  # NOTE: test perf
-
     rank = args.gpus_per_node * node + local_rank
 
+    # Need randomness across processes, for sampler, augmentation, noise etc.
+    # Note DDP broadcasts initial model states from rank 0
+    torch.manual_seed(args.seed + rank)
+    #torch.backends.cudnn.deterministic = True  # NOTE: test perf
+
     dist_init(rank, args)
 
     train_dataset = FieldDataset(
@@ -72,7 +74,13 @@ def gpu_worker(local_rank, node, args):
         rank=rank,
         world_size=args.world_size,
     )
-    if not args.div_data:
+    if args.div_data:
+        train_sampler = GroupedRandomSampler(
+            train_dataset,
+            group_size=None if args.cache_maxsize is None else
+                       args.cache_maxsize * train_dataset.ncrop,
+        )
+    else:
         try:
             train_sampler = DistributedSampler(train_dataset, shuffle=True)
         except TypeError:
@@ -80,8 +88,8 @@ def gpu_worker(local_rank, node, args):
     train_loader = DataLoader(
         train_dataset,
         batch_size=args.batches,
-        shuffle=args.div_data,
+        shuffle=False,
-        sampler=None if args.div_data else train_sampler,
+        sampler=train_sampler,
         num_workers=args.loader_workers,
         pin_memory=True,
     )
@@ -104,7 +112,9 @@ def gpu_worker(local_rank, node, args):
             rank=rank,
             world_size=args.world_size,
         )
-        if not args.div_data:
+        if args.div_data:
+            val_sampler = None
+        else:
             try:
                 val_sampler = DistributedSampler(val_dataset, shuffle=False)
             except TypeError:
@@ -113,7 +123,7 @@ def gpu_worker(local_rank, node, args):
             val_dataset,
             batch_size=args.batches,
             shuffle=False,
-            sampler=None if args.div_data else val_sampler,
+            sampler=val_sampler,
             num_workers=args.loader_workers,
             pin_memory=True,
         )
@@ -278,6 +288,12 @@ def gpu_worker(local_rank, node, args):
             os.symlink(state_file, tmp_link)  # workaround to overwrite
             os.rename(tmp_link, ckpt_link)
 
+    if args.cache:
+        print('rank {} train data: {}'.format(
+            rank, train_dataset.get_fields.cache_info()))
+        print('rank {} val   data: {}'.format(
+            rank, val_dataset.get_fields.cache_info()))
+
     dist.destroy_process_group()