From 8a6db0b399706798caf8900a4595434ce59ee07e Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Sun, 14 Apr 2024 11:37:16 +0000 Subject: [PATCH] refactor dynamic training --- .gitignore | 5 +- opensora/dataset/t2v_datasets.py | 83 +++++++++++-------- opensora/models/diffusion/latte/modules.py | 18 ++-- opensora/utils/dataset_utils.py | 54 +++++++----- scripts/accelerate_configs/ddp_config.yaml | 6 +- .../train_videoae_65x512x512.sh | 8 +- 6 files changed, 103 insertions(+), 71 deletions(-) diff --git a/.gitignore b/.gitignore index 960f44d7a..650452d34 100644 --- a/.gitignore +++ b/.gitignore @@ -18,4 +18,7 @@ wandb/ *.png *.gif *.pth -*.pt \ No newline at end of file +*.pt +cache_dir/ +wandb/ +test/ \ No newline at end of file diff --git a/opensora/dataset/t2v_datasets.py b/opensora/dataset/t2v_datasets.py index 6741d87e1..e979aadac 100644 --- a/opensora/dataset/t2v_datasets.py +++ b/opensora/dataset/t2v_datasets.py @@ -14,6 +14,10 @@ from opensora.utils.dataset_utils import DecordInit from opensora.utils.utils import text_preprocessing +def random_video_noise(t, c, h, w): + vid = torch.rand(t, c, h, w) * 255.0 + vid = vid.to(torch.uint8) + return vid class T2V_dataset(Dataset): def __init__(self, args, transform, temporal_sample, tokenizer): @@ -40,47 +44,58 @@ class T2V_dataset(Dataset): def __getitem__(self, idx): try: - # video = torch.randn(3, 16, 128, 128) - # input_ids = torch.ones(1, 120).to(torch.long).squeeze(0) - # cond_mask = torch.cat([torch.ones(1, 60).to(torch.long), torch.ones(1, 60).to(torch.long)], dim=1).squeeze(0) - # return video, input_ids, cond_mask - video_path = self.samples[idx]['path'] - video = self.decord_read(video_path) - video = self.transform(video) # T C H W -> T C H W - video = video.transpose(0, 1) # T C H W -> C T H W - text = self.samples[idx]['cap'][0] - - text = text_preprocessing(text) - text_tokens_and_mask = self.tokenizer( - text, - max_length=self.model_max_length, - padding='max_length', - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors='pt' - ) - input_ids = text_tokens_and_mask['input_ids'].squeeze(0) - cond_mask = text_tokens_and_mask['attention_mask'].squeeze(0) - + video_data = self.get_video(idx) + image_data = {} if self.use_image_num != 0 and self.use_img_from_vid: - select_image_idx = np.linspace(0, self.num_frames-1, self.use_image_num, dtype=int) - assert self.num_frames >= self.use_image_num - images = video[:, select_image_idx] # c, num_img, h, w - video = torch.cat([video, images], dim=1) # c, num_frame+num_img, h, w - input_ids = torch.stack([input_ids] * (1+self.use_image_num)) # 1+self.use_image_num, l - cond_mask = torch.stack([cond_mask] * (1+self.use_image_num)) # 1+self.use_image_num, l + image_data = self.get_image_from_video(video_data) elif self.use_image_num != 0 and not self.use_img_from_vid: - images, captions = self.img_cap_list[idx] - raise NotImplementedError + image_data = self.get_image(idx) else: - pass - - return video, input_ids, cond_mask + raise NotImplementedError + return dict(video_data=video_data, image_data=image_data) except Exception as e: print(f'Error with {e}, {self.samples[idx]}') return self.__getitem__(random.randint(0, self.__len__() - 1)) + def get_video(self, idx): + # video = random.choice([random_video_noise(65, 3, 720, 360) * 255, random_video_noise(65, 3, 1024, 1024), random_video_noise(65, 3, 360, 720)]) + # print('random shape', video.shape) + # input_ids = torch.ones(1, 120).to(torch.long).squeeze(0) + # cond_mask = torch.cat([torch.ones(1, 60).to(torch.long), torch.ones(1, 60).to(torch.long)], dim=1).squeeze(0) + + video_path = self.samples[idx]['path'].replace('/remote-home1/dataset/data_split_tt', '/yzwldata01/dataset/split_1024') + video = self.decord_read(video_path) + video = self.transform(video) # T C H W -> T C H W + video = video.transpose(0, 1) # T C H W -> C T H W + text = self.samples[idx]['cap'][0] + + text = text_preprocessing(text) + text_tokens_and_mask = self.tokenizer( + text, + max_length=self.model_max_length, + padding='max_length', + truncation=True, + return_attention_mask=True, + add_special_tokens=True, + return_tensors='pt' + ) + input_ids = text_tokens_and_mask['input_ids'] + cond_mask = text_tokens_and_mask['attention_mask'] + return dict(video=video, input_ids=input_ids, cond_mask=cond_mask) + + def get_image_from_video(self, video_data): + select_image_idx = np.linspace(0, self.num_frames-1, self.use_image_num, dtype=int) + assert self.num_frames >= self.use_image_num + image = [video_data['video'][:, i:i+1] for i in select_image_idx] # num_img [c, 1, h, w] + input_ids = video_data['input_ids'].repeat(self.use_image_num, 1) # self.use_image_num, l + cond_mask = video_data['cond_mask'].repeat(self.use_image_num, 1) # self.use_image_num, l + return dict(image=image, input_ids=input_ids, cond_mask=cond_mask) + + def get_image(self, idx): + raise NotImplementedError + image_data = self.img_cap_list[idx] + return dict(image=image, input_ids=input_ids, cond_mask=cond_mask) + def tv_read(self, path): vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW') total_frames = len(vframes) diff --git a/opensora/models/diffusion/latte/modules.py b/opensora/models/diffusion/latte/modules.py index 2ad20343a..63fa7aa0b 100644 --- a/opensora/models/diffusion/latte/modules.py +++ b/opensora/models/diffusion/latte/modules.py @@ -263,10 +263,6 @@ class Attention(nn.Module): self.dropout = dropout self.use_rope = use_rope - if self.use_rope: - self.rope2d = RoPE2D(freq=100) - self.rope1d = RoPE1D(freq=100) - # we make use of this private variable to know whether this class is loaded # with an deprecated state dict so that we can convert it on the fly self._from_deprecated_attn_block = _from_deprecated_attn_block @@ -350,7 +346,7 @@ class Attention(nn.Module): # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1 if processor is None: processor = ( - AttnProcessor2_0(attention_mode) if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor() + AttnProcessor2_0(attention_mode, use_rope) if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor() ) self.set_processor(processor) @@ -843,8 +839,13 @@ class AttnProcessor2_0: Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). """ - def __init__(self, attention_mode='xformers'): + def __init__(self, attention_mode='xformers', use_rope=False): self.attention_mode = attention_mode + self.use_rope = use_rope + + if self.use_rope: + self.rope2d = RoPE2D(freq=100) + self.rope1d = RoPE1D(freq=100) if not hasattr(F, "scaled_dot_product_attention"): raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") @@ -904,7 +905,10 @@ class AttnProcessor2_0: key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - if self.rope is not None: + if self.use_rope: + print('use rope') + print('use rope') + print('use rope') # require the shape of (batch_size x nheads x ntokens x dim) if position_q.ndim == 3: query = self.rope2d(query, position_q) diff --git a/opensora/utils/dataset_utils.py b/opensora/utils/dataset_utils.py index 1883674a0..dc0c3656c 100644 --- a/opensora/utils/dataset_utils.py +++ b/opensora/utils/dataset_utils.py @@ -1,5 +1,5 @@ import math - +from einops import rearrange import decord from torch.nn import functional as F import torch @@ -58,45 +58,55 @@ class Collate: self.num_frames = args.num_frames self.use_image_num = args.use_image_num self.max_thw = (self.num_frames, self.max_image_size, self.max_image_size) - self.max_ihw = (self.use_image_num, self.max_image_size, self.max_image_size) + self.max_1hw = (1, self.max_image_size, self.max_image_size) + def package(self, batch): + batch_tubes_vid = [i['video_data']['video'] for i in batch] # b [c t h w] + input_ids_vid = torch.stack([i['video_data']['input_ids'] for i in batch]) # b 1 l + cond_mask_vid = torch.stack([i['video_data']['cond_mask'] for i in batch]) # b 1 l + batch_tubes_img, input_ids_img, cond_mask_img = None, None, None + if self.use_image_num != 0: + import ipdb;ipdb.set_trace() + batch_tubes_img = [j for i in batch for j in i['image_data']['image']] # b*num_img [c 1 h w] + input_ids_img = torch.stack([i['image_data']['input_ids'] for i in batch]) # b image_num l + cond_mask_img = torch.stack([i['image_data']['cond_mask'] for i in batch]) # b image_num l + return batch_tubes_vid, input_ids_vid, cond_mask_vid, batch_tubes_img, input_ids_img, cond_mask_img def __call__(self, batch): - unzip = tuple(zip(*batch)) - batch_tubes, input_ids, cond_mask = unzip - input_ids = torch.stack(input_ids).squeeze(1) - cond_mask = torch.stack(cond_mask).squeeze(1) + batch_tubes_vid, input_ids_vid, cond_mask_vid, batch_tubes_img, input_ids_img, cond_mask_img = self.package(batch) ds_stride = self.ae_stride * self.patch_size t_ds_stride = self.ae_stride_t * self.patch_size_t if self.use_image_num == 0: - pad_batch_tubes, attention_mask = process(batch_tubes, t_ds_stride, ds_stride, self.max_thw, self.ae_stride_thw, self.patch_size_thw) + pad_batch_tubes, attention_mask = process(batch_tubes_vid, t_ds_stride, ds_stride, + self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True) # attention_mask: b t hw, use first frame attention_mask = attention_mask[:, 0] # b hw + input_ids, cond_mask = input_ids_vid.squeeze(1), cond_mask_vid.squeeze(1) # b 1 l -> b l else: - batch_tubes_vid = [i[:, :-self.use_image_num] for i in batch_tubes] - pad_batch_tubes_vid, attention_mask_vid = process(batch_tubes_vid, t_ds_stride, ds_stride, self.max_thw, self.ae_stride_thw, self.patch_size_thw) - print('pad_batch_tubes_vid.shape, attention_mask_vid.shape', pad_batch_tubes_vid.shape, attention_mask_vid.shape) + pad_batch_tubes_vid, attention_mask_vid = self.process(batch_tubes_vid, t_ds_stride, ds_stride, + self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True) # attention_mask_vid: b t hw, use first frame attention_mask_vid = attention_mask_vid[:, :1] # b 1 hw - batch_tubes_img = [i[:, -self.use_image_num:] for i in batch_tubes] - pad_batch_tubes_img, attention_mask_img = process(batch_tubes_img, 1, ds_stride, self.max_ihw, self.ae_stride_1hw, self.patch_size_1hw) - print('pad_batch_tubes_img.shape, attention_mask_img.shape', pad_batch_tubes_img.shape, attention_mask_img.shape) - pad_batch_tubes = torch.cat([pad_batch_tubes_vid, pad_batch_tubes_img], dim=3) + pad_batch_tubes_img, attention_mask_img = self.process(batch_tubes_img, 1, ds_stride, + self.max_1hw, self.ae_stride_1hw, self.patch_size_1hw, extra_1=False) + pad_batch_tubes_img = rearrange(pad_batch_tubes_img, '(b i) c 1 h w -> b c i h w', i=self.use_image_num) + attention_mask_img = rearrange(attention_mask_img, '(b i) 1 hw -> b i hw', i=self.use_image_num) + pad_batch_tubes = torch.cat([pad_batch_tubes_vid, pad_batch_tubes_img], dim=2) # concat at temporal, video first # attention_mask_img: b num_img hw attention_mask = torch.cat([attention_mask_vid, attention_mask_img], dim=1) # b 1+num_img hw + input_ids = torch.cat([input_ids_vid, input_ids_img], dim=1) # b 1+num_img hw + cond_mask = torch.cat([cond_mask_vid, cond_mask_img], dim=1) # b 1+num_img hw return pad_batch_tubes, attention_mask, input_ids, cond_mask - def process(batch_tubes, t_ds_stride, ds_stride, max_thw, ae_stride_thw, patch_size_thw): + def process(self, batch_tubes, t_ds_stride, ds_stride, max_thw, ae_stride_thw, patch_size_thw, extra_1): # pad to max multiple of ds_stride batch_input_size = [i.shape for i in batch_tubes] # [(c t h w), (c t h w)] max_t, max_h, max_w = max_thw - # max_t, max_h, max_w = max([i[1] for i in batch_input_size]), \ - # max([i[2] for i in batch_input_size]), \ - # max([i[3] for i in batch_input_size]) - pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t, t_ds_stride), \ + pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t-1 if extra_1 else max_t, t_ds_stride), \ pad_to_multiple(max_h, ds_stride), \ pad_to_multiple(max_w, ds_stride) + pad_max_t = pad_max_t + 1 if extra_1 else pad_max_t each_pad_t_h_w = [[pad_max_t - i.shape[1], pad_max_h - i.shape[2], pad_max_w - i.shape[3]] for i in batch_tubes] @@ -108,13 +118,13 @@ class Collate: # make attention_mask max_tube_size = [pad_max_t, pad_max_h, pad_max_w] - max_latent_size = [max_tube_size[0] // ae_stride_thw[0], + max_latent_size = [((max_tube_size[0]-1) // ae_stride_thw[0] + 1) if extra_1 else (max_tube_size[0] // ae_stride_thw[0]), max_tube_size[1] // ae_stride_thw[1], max_tube_size[2] // ae_stride_thw[2]] - max_patchify_latent_size = [max_latent_size[0] // patch_size_thw[0], + max_patchify_latent_size = [((max_latent_size[0]-1) // patch_size_thw[0] + 1) if extra_1 else (max_latent_size[0] // patch_size_thw[0]), max_latent_size[1] // patch_size_thw[1], max_latent_size[2] // patch_size_thw[2]] - valid_patchify_latent_size = [[int(math.ceil(i[1] / t_ds_stride)), + valid_patchify_latent_size = [[int(math.ceil((i[1]-1) / t_ds_stride)) + 1 if extra_1 else int(math.ceil(i[1] / t_ds_stride)), int(math.ceil(i[2] / ds_stride)), int(math.ceil(i[3] / ds_stride))] for i in batch_input_size] attention_mask = [F.pad(torch.ones(i), diff --git a/scripts/accelerate_configs/ddp_config.yaml b/scripts/accelerate_configs/ddp_config.yaml index fa6cfbdd6..f1481bd2e 100644 --- a/scripts/accelerate_configs/ddp_config.yaml +++ b/scripts/accelerate_configs/ddp_config.yaml @@ -1,11 +1,11 @@ compute_environment: LOCAL_MACHINE -distributed_type: MULTI_GPU +distributed_type: NO fsdp_config: {} machine_rank: 0 main_process_ip: null main_process_port: 29501 main_training_function: main num_machines: 1 -num_processes: 8 -gpu_ids: 0,1,2,3,4,5,6,7 +num_processes: 1 +gpu_ids: 0, use_cpu: false \ No newline at end of file diff --git a/scripts/text_condition/train_videoae_65x512x512.sh b/scripts/text_condition/train_videoae_65x512x512.sh index d3a9729bd..9511ac6b6 100644 --- a/scripts/text_condition/train_videoae_65x512x512.sh +++ b/scripts/text_condition/train_videoae_65x512x512.sh @@ -2,7 +2,7 @@ export WANDB_KEY="" export ENTITY="" export PROJECT="test" accelerate launch \ - --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ + --config_file scripts/accelerate_configs/ddp_config.yaml \ opensora/train/train_t2v.py \ --model LatteT2V-XL/122 \ --text_encoder_name DeepFloyd/t5-v1_1-xxl \ @@ -17,15 +17,15 @@ accelerate launch \ --max_image_size 512 \ --gradient_checkpointing \ --attention_mode xformers \ - --train_batch_size=2 \ - --dataloader_num_workers 10 \ + --train_batch_size=4 \ + --dataloader_num_workers 0 \ --gradient_accumulation_steps=1 \ --max_train_steps=1000000 \ --learning_rate=2e-05 \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ --mixed_precision="bf16" \ - --report_to="wandb" \ + --report_to="tensorboard" \ --checkpointing_steps=500 \ --output_dir="test" \ --allow_tf32 \