Update training module name to open_clip_train

src/open_clip_train/main.py

@@ -29,13 +29,13 @@
     hvd = None
 
 from open_clip import create_model_and_transforms, trace_model, get_tokenizer, create_loss
-from training.data import get_data
-from training.distributed import is_master, init_distributed_device, broadcast_object
-from training.logger import setup_logging
-from training.params import parse_args
-from training.scheduler import cosine_lr, const_lr, const_lr_cooldown
-from training.train import train_one_epoch, evaluate
-from training.file_utils import pt_load, check_exists, start_sync_process, remote_sync
+from open_clip_train.data import get_data
+from open_clip_train.distributed import is_master, init_distributed_device, broadcast_object
+from open_clip_train.logger import setup_logging
+from open_clip_train.params import parse_args
+from open_clip_train.scheduler import cosine_lr, const_lr, const_lr_cooldown
+from open_clip_train.train import train_one_epoch, evaluate
+from open_clip_train.file_utils import pt_load, check_exists, start_sync_process, remote_sync
 
 
 LATEST_CHECKPOINT_NAME = "epoch_latest.pt"

src/open_clip_train/train.py

@@ -15,9 +15,9 @@
     wandb = None
 
 from open_clip import get_input_dtype, CLIP, CustomTextCLIP
-from .distributed import is_master
-from .zero_shot import zero_shot_eval
-from .precision import get_autocast
+from open_clip_train.distributed import is_master
+from open_clip_train.zero_shot import zero_shot_eval
+from open_clip_train.precision import get_autocast
 
 
 class AverageMeter(object):

src/open_clip_train/zero_shot.py

@@ -5,7 +5,7 @@
 
 from open_clip import get_input_dtype, get_tokenizer, build_zero_shot_classifier, \
     IMAGENET_CLASSNAMES, OPENAI_IMAGENET_TEMPLATES
-from .precision import get_autocast
+from open_clip_train.precision import get_autocast
 
 
 def accuracy(output, target, topk=(1,)):

tests/test_hf_model.py

@@ -4,6 +4,7 @@
 from open_clip.hf_model import _POOLERS, HFTextEncoder
 from transformers import AutoConfig
 from transformers.modeling_outputs import BaseModelOutput
+
 # test poolers
 def test_poolers():
     bs, sl, d = 2, 10, 5

tests/test_num_shards.py

@@ -1,6 +1,6 @@
 import pytest
 
-from training.data import get_dataset_size
+from open_clip_train.data import get_dataset_size
 
 @pytest.mark.parametrize(
     "shards,expected_size",

tests/test_training_simple.py

@@ -2,9 +2,8 @@
 import os
 import sys
 import pytest
-from PIL import Image
 import torch
-from training.main import main
+from open_clip_train.main import main
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 

tests/test_wds.py

@@ -6,9 +6,9 @@
 import io
 from PIL import Image
 
-from training.data import get_wds_dataset
-from training.params import parse_args
-from training.main import random_seed
+from open_clip_train.data import get_wds_dataset
+from open_clip_train.params import parse_args
+from open_clip_train.main import random_seed
 
 TRAIN_NUM_SAMPLES = 10_000
 RTOL = 0.2