萤火跑模型 | CLIP 模型

# image_encoder - ResNet or Vision Transformer
# text_encoder  - CBOW or Text Transformer
# I[n, h, w, c] - minibatch of aligned images
# T[n, l]       - minibatch of aligned texts
# W_i[d_i, d_e] - learned proj of image to embed
# W_t[d_t, d_e] - learned proj of text to embed
# t             - learned temperature parameter
# extract feature representations of each modality
I_f = image_encoder(I) #[n, d_i]
T_f = text_encoder(T)  #[n, d_t]
# joint multimodal embedding [n, d_e]
I_e = l2_normalize(np.dot(I_f, W_i), axis=1)
T_e = l2_normalize(np.dot(T_f, W_t), axis=1)
# scaled pairwise cosine similarities [n, n]
logits = np.dot(I_e, T_e.T) * np.exp(t)
# symmetric loss function
labels = np.arange(n)
loss_i = cross_entropy_loss(logits, labels, axis=0)
loss_t = cross_entropy_loss(logits, labels, axis=1)
loss   = (loss_i + loss_t)/2

title	filepath
a man on a scooter , taking part in the rally	http://l7.alamy.com/zooms/5e84d0b73a544fe8a8e5ac0d20052685/a-man-on-a-lambretta-scooter-taking-part-in-the-daily-express-rally-ddmmff.jpg
the dentist drill the tooth with a turbine	http://l7.alamy.com/zooms/da764b930a6e477da792ddb8b83b1ecb/the-dentist-drill-the-tooth-with-a-turbine-hn3aph.jpg
...

title	filepath
a man on a scooter , taking part in the rally	/demo/path/clip/cc_data_1/val/29/29.jpg
the dentist drill the tooth with a turbine	/demo/path/clip/cc_data_1/val/19/19.jpg
...

'--train-ddr-perfix', '/3fs-jd/prod/private/nxt/clip/cc_data_1_proc/train',
'--val-ddr-perfix', '/3fs-jd/prod/private/nxt/clip/cc_data_1_proc/val',

import haienv
haienv.set_env("202111")

import os
import sys
sys.path.append(f'{os.getcwd()}/src')

class FFDataset(fftorch.Dataset):
    def __init__(self, input_filename, transforms, img_key, caption_key, sep="\t", ddr_perfix=None):
        logging.info(f'[FFDataset] Loading csv data from {input_filename}. use ffrecord: {ddr_perfix}')
        df = pd.read_csv(input_filename, sep=sep)

        self.images = df[img_key].tolist()
        self.captions = df[caption_key].tolist()
        self.transforms = transforms
        self.ddr_perfix = ddr_perfix
        super().__init__(self.ddr_perfix, check_data=True)
        logging.info('[FFDataset] Done loading data.')

    def __len__(self):
        return len(self.captions)

    def process(self, indexes, data):
        '''
        在每次读取样本数据的时候，
        FFDataset 先从 FFRecord 文件中读取样本的二进制数据，
        然后将原始二进制数据传递给 process 函数。
        用户需要继承 ffrecord.torch.Dataset 并自定义 process 函数以对原始二进制数据进行处理。
        '''
        samples = []

        try:
            for bytes_ in data:
                sample = deserialize(bytes_)
                image_from_bytes = sample['image_bytes']
                images = self.transforms(image_from_bytes)
                texts = tokenize([str(sample['caption'])])[0]
        
                samples.append((images, texts))
        except:
            raise

        return samples

# main 函数里面处理 whole_life_state:
whole_life_state = hfai.get_whole_life_state()
print('whole_life_state:', whole_life_state)

if whole_life_state == 1:
    # 只要不是第一次运行，我们就尝试加载最近的 checkpoint：
    args.resume = os.path.join(args.checkpoint_path, f"epoch_latest.pt")

hfai.set_whole_life_state(1)

# 加载 checkpoint:
if args.resume is not None:
    if os.path.isfile(args.resume):
        checkpoint = torch.load(args.resume, map_location=loc)
        start_epoch = checkpoint["epoch"]
        sd = checkpoint["state_dict"]
        model.load_state_dict(sd)
        if optimizer is not None:
            optimizer.load_state_dict(checkpoint["optimizer"])
        logging.info(
            f"=> loaded checkpoint resume:'{args.resume}' (epoch {checkpoint['epoch']})"
        )
    # ...

2021-12-22,22:34:14 | INFO | Rank 0 | Use GPU: 0 for training
2021-12-22,22:34:14 | INFO | Rank 4 | Use GPU: 4 for training
2021-12-22,22:34:14 | INFO | Rank 6 | Use GPU: 6 for training
2021-12-22,22:34:14 | INFO | Rank 3 | Use GPU: 3 for training
2021-12-22,22:34:14 | INFO | Rank 7 | Use GPU: 7 for training
2021-12-22,22:34:14 | INFO | Rank 1 | Use GPU: 1 for training
2021-12-22,22:34:14 | INFO | Rank 2 | Use GPU: 2 for training
2021-12-22,22:34:14 | INFO | Rank 5 | Use GPU: 5 for training
... ...
Train Epoch: 0 [0/2850879 (0%)]	Loss: 8.067907	Data (t) 42.464	Batch (t) 72.875	LR: 0.000000	logit_scale 2.659
Train Epoch: 0 [3840/2850879 (1%)]	Loss: 8.055143	Data (t) 0.008	Batch (t) 0.356	LR: 0.000001	logit_scale 2.659