[Minor] Reduce memory usage

Author: zsyzzsoft

DiffAugment-stylegan2/run_few_shot.py

@@ -42,7 +42,7 @@ def run(dataset, resolution, result_dir, DiffAugment, num_gpus, batch_size, tota
     metric_args = EasyDict(cache_dir=dataset, num_repeats=num_repeats)
 
     desc = 'DiffAugment-stylegan2' if DiffAugment else 'stylegan2'
-    dataset_args = EasyDict(tfrecord_dir=dataset, resolution=resolution, shuffle_mb=256, prefetch_mb=256, from_tfrecords=True)
+    dataset_args = EasyDict(tfrecord_dir=dataset, resolution=resolution, from_tfrecords=True)
     desc += '-' + os.path.basename(dataset)
     if resolution is not None:
         desc += '-{}'.format(resolution)

DiffAugment-stylegan2/training/dataset.py

@@ -93,7 +93,7 @@ def __init__(self,
 
             dset = tf.data.TFRecordDataset(tfr_file, compression_type='', buffer_size=buffer_mb << 20)
 
-            self._np_labels = np.zeros([1 << 30], dtype=np.int32)
+            self._np_labels = np.zeros([1 << 25], dtype=np.int32)
             if self.label_file is not None and max_label_size != 0:
                 self._np_labels = np.load(self.label_file).astype(np.int32)
                 self.label_size = self._np_labels.max() + 1