rename inmem to gpu_trans. remoce one_iter.

facebookresearch · Sep 25, 2024 · ee4664b · ee4664b
1 parent 5aaafd9
commit ee4664b
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 8 deletions.
diff --git a/config/metaclip_2_5b.py b/config/metaclip_2_5b.py
@@ -11,8 +11,7 @@
 
 @dataclass
 class b32_fullcc(Config):
-    one_iter=True
-    inmem=True
+    gpu_trans=True
     engine="train_one_epoch_ex"
     eval_steps=5000
     save_frequency=1

diff --git a/config/metaclip_400m.py b/config/metaclip_400m.py
@@ -11,7 +11,7 @@
 
 @dataclass
 class b32_400m(Config):
-    inmem=True
+    gpu_trans=True
     engine="train_one_epoch_ex"
     eval_steps=5000
     save_frequency=1

diff --git a/src/open_clip/factory.py b/src/open_clip/factory.py
@@ -137,7 +137,7 @@ def create_model_and_transforms(
     pretrained_image: bool = False,
     mean: Optional[Tuple[float, ...]] = None,
     std: Optional[Tuple[float, ...]] = None,
-    inmem = False,
+    gpu_trans = False,
     clip_model: str = "CLIP",
 ):
     model = create_model(
@@ -146,7 +146,7 @@ def create_model_and_transforms(
         pretrained_image=pretrained_image,
         clip_model=clip_model,
     )
-    preprocess_train = image_transform(model.visual.image_size, is_train=True, mean=mean, std=std, inmem=inmem)
+    preprocess_train = image_transform(model.visual.image_size, is_train=True, mean=mean, std=std, gpu_trans=gpu_trans)
     preprocess_val = image_transform(model.visual.image_size, is_train=False, mean=mean, std=std)
     return model, preprocess_train, preprocess_val
 

diff --git a/src/open_clip/transform.py b/src/open_clip/transform.py
@@ -54,7 +54,7 @@ def image_transform(
     std: Optional[Tuple[float, ...]] = None,
     resize_longest_max: bool = False,
     fill_color: int = 0,
-    inmem = False
+    gpu_trans = False
 ):
     mean = mean or (0.48145466, 0.4578275, 0.40821073)  # OpenAI dataset mean
     std = std or (0.26862954, 0.26130258, 0.27577711)  # OpenAI dataset std
@@ -64,7 +64,7 @@ def image_transform(
 
     normalize = Normalize(mean=mean, std=std)
     if is_train:
-        if inmem:
+        if gpu_trans:
             return Compose([
                 RandomResizedCrop(image_size, scale=(0.9, 1.0), interpolation=InterpolationMode.BICUBIC),
                 _convert_to_rgb,

diff --git a/src/training/main.py b/src/training/main.py
@@ -115,7 +115,7 @@ def main(args):
         force_quick_gelu=args.force_quick_gelu,
         pretrained_image=args.pretrained_image,
         mean=mean, std=std,
-        inmem=hasattr(args, "inmem"),
+        gpu_trans=hasattr(args, "gpu_trans") and args.gpu_trans,
         clip_model=args.clip_model,
     )
     random_seed(args.seed, args.rank)