522 lines
21 KiB
Python
522 lines
21 KiB
Python
#!/usr/bin/env python
|
|
# coding=utf-8
|
|
# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
|
|
import copy
|
|
import logging
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import diffusers
|
|
import torch
|
|
import torch.utils.checkpoint
|
|
import transformers
|
|
import yaml
|
|
from accelerate import Accelerator
|
|
from accelerate.utils import DeepSpeedPlugin, ProjectConfiguration, set_seed
|
|
from diffusers.optimization import get_scheduler
|
|
from diffusers.utils import is_wandb_available
|
|
from huggingface_hub import create_repo, upload_folder
|
|
from tqdm.auto import tqdm
|
|
from safetensors.torch import load_model
|
|
|
|
from models.ema_model import EMAModel
|
|
from models.multimodal_encoder.siglip_encoder import SiglipVisionTower
|
|
from models.multimodal_encoder.t5_encoder import T5Embedder
|
|
from models.rdt_runner import RDTRunner
|
|
from train.dataset import DataCollatorForVLAConsumerDataset, VLAConsumerDataset
|
|
from train.sample import log_sample_res
|
|
|
|
if is_wandb_available():
|
|
import wandb
|
|
|
|
|
|
def save_model_card(repo_id: str, base_model=str, repo_folder=None):
|
|
yaml = f"""
|
|
---
|
|
license: mit
|
|
base_model: {base_model}
|
|
language:
|
|
- en
|
|
pipeline_tag: robotics
|
|
library_name: transformers
|
|
tags:
|
|
- robotics
|
|
- pytorch
|
|
- multimodal
|
|
- pretraining
|
|
- vla
|
|
- diffusion
|
|
- rdt
|
|
---
|
|
"""
|
|
model_card = f"""
|
|
# RDT - {repo_id}
|
|
|
|
This is a RDT model derived from {base_model}. The weights were trained using [RDT](https://rdt-robotics.github.io/rdt-robotics/).
|
|
"""
|
|
with open(os.path.join(repo_folder, "README.md"), "w") as f:
|
|
f.write(yaml + model_card)
|
|
|
|
|
|
def train(args, logger):
|
|
# Read the config
|
|
with open(args.config_path, "r") as fp:
|
|
config = yaml.safe_load(fp)
|
|
|
|
with open(args.model_config_path, "r") as f:
|
|
model_config = yaml.safe_load(f)
|
|
# print(model_config)
|
|
args.output_dir = model_config["checkpoint_path"]
|
|
logging_dir = Path(args.output_dir, args.logging_dir)
|
|
|
|
accelerator_project_config = ProjectConfiguration(total_limit=args.checkpoints_total_limit)
|
|
accelerator = Accelerator(
|
|
deepspeed_plugin=(DeepSpeedPlugin(hf_ds_config=args.deepspeed) if args.deepspeed is not None else None),
|
|
gradient_accumulation_steps=args.gradient_accumulation_steps,
|
|
mixed_precision=args.mixed_precision,
|
|
log_with=args.report_to,
|
|
project_dir=logging_dir,
|
|
project_config=accelerator_project_config,
|
|
)
|
|
|
|
if args.report_to == "wandb":
|
|
if not is_wandb_available():
|
|
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
|
|
|
|
# Make one log on every process with the configuration for debugging.
|
|
logging.basicConfig(
|
|
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
|
datefmt="%m/%d/%Y %H:%M:%S",
|
|
level=logging.INFO,
|
|
filename=args.output_log_path,
|
|
filemode='w',
|
|
)
|
|
logger.info(accelerator.state, main_process_only=False)
|
|
if accelerator.is_local_main_process:
|
|
transformers.utils.logging.set_verbosity_warning()
|
|
diffusers.utils.logging.set_verbosity_info()
|
|
else:
|
|
transformers.utils.logging.set_verbosity_error()
|
|
diffusers.utils.logging.set_verbosity_error()
|
|
|
|
# If passed along, set the training seed now.
|
|
if args.seed is not None:
|
|
set_seed(args.seed)
|
|
|
|
# Handle the repository creation
|
|
if accelerator.is_main_process:
|
|
if args.output_dir is not None:
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
|
|
if args.push_to_hub:
|
|
repo_id = create_repo(
|
|
repo_id=args.hub_model_id or Path(args.output_dir).name,
|
|
exist_ok=True,
|
|
token=args.hub_token,
|
|
).repo_id
|
|
|
|
# For mixed precision training we cast the text_encoder and vae weights to half-precision
|
|
# as these models are only used for inference, keeping weights in full precision is not required.
|
|
weight_dtype = torch.float32
|
|
if accelerator.mixed_precision == "fp16":
|
|
weight_dtype = torch.float16
|
|
elif accelerator.mixed_precision == "bf16":
|
|
weight_dtype = torch.bfloat16
|
|
|
|
if args.precomp_lang_embed:
|
|
tokenizer, text_encoder = None, None
|
|
else:
|
|
text_embedder = T5Embedder(
|
|
from_pretrained=args.pretrained_text_encoder_name_or_path,
|
|
model_max_length=config["dataset"]["tokenizer_max_length"],
|
|
device=accelerator.device,
|
|
)
|
|
tokenizer, text_encoder = text_embedder.tokenizer, text_embedder.model
|
|
|
|
vision_encoder = SiglipVisionTower(vision_tower=args.pretrained_vision_encoder_name_or_path, args=None)
|
|
image_processor = vision_encoder.image_processor
|
|
|
|
# Load from a pretrained checkpoint
|
|
if args.pretrained_model_name_or_path is not None and not os.path.isfile(args.pretrained_model_name_or_path):
|
|
logger.info("Constructing model from pretrained checkpoint.")
|
|
rdt = RDTRunner.from_pretrained(args.pretrained_model_name_or_path)
|
|
else:
|
|
logger.info("Constructing model from provided config.")
|
|
# Calculate the image condition length
|
|
img_cond_len = (config["common"]["img_history_size"] * config["common"]["num_cameras"] *
|
|
vision_encoder.num_patches)
|
|
rdt = RDTRunner(
|
|
action_dim=config["common"]["state_dim"],
|
|
pred_horizon=config["common"]["action_chunk_size"],
|
|
config=config["model"],
|
|
lang_token_dim=config["model"]["lang_token_dim"],
|
|
img_token_dim=config["model"]["img_token_dim"],
|
|
state_token_dim=config["model"]["state_token_dim"],
|
|
max_lang_cond_len=config["dataset"]["tokenizer_max_length"],
|
|
img_cond_len=img_cond_len,
|
|
img_pos_embed_config=[
|
|
# No initial pos embed in the last grid size
|
|
# since we've already done in ViT
|
|
(
|
|
"image",
|
|
(
|
|
config["common"]["img_history_size"],
|
|
config["common"]["num_cameras"],
|
|
-vision_encoder.num_patches,
|
|
),
|
|
),
|
|
],
|
|
lang_pos_embed_config=[
|
|
# Similarly, no initial pos embed for language
|
|
("lang", -config["dataset"]["tokenizer_max_length"]),
|
|
],
|
|
dtype=weight_dtype,
|
|
)
|
|
|
|
ema_rdt = copy.deepcopy(rdt)
|
|
ema_model = EMAModel(
|
|
ema_rdt,
|
|
update_after_step=config["model"]["ema"]["update_after_step"],
|
|
inv_gamma=config["model"]["ema"]["inv_gamma"],
|
|
power=config["model"]["ema"]["power"],
|
|
min_value=config["model"]["ema"]["min_value"],
|
|
max_value=config["model"]["ema"]["max_value"],
|
|
)
|
|
|
|
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
|
|
# which ensure saving model in huggingface format (config.json + pytorch_model.bin)
|
|
def save_model_hook(models, weights, output_dir):
|
|
if accelerator.is_main_process:
|
|
for model in models:
|
|
model_to_save = model.module if hasattr(model, "module") else model # type: ignore
|
|
if isinstance(model_to_save, type(accelerator.unwrap_model(rdt))):
|
|
model_to_save.save_pretrained(output_dir)
|
|
|
|
accelerator.register_save_state_pre_hook(save_model_hook)
|
|
|
|
if args.gradient_checkpointing:
|
|
# TODO:
|
|
raise NotImplementedError("Gradient checkpointing is not yet implemented.")
|
|
|
|
# Enable TF32 for faster training on Ampere GPUs,
|
|
# cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
|
|
if args.allow_tf32:
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
|
|
if args.scale_lr:
|
|
args.learning_rate = (args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size *
|
|
accelerator.num_processes)
|
|
|
|
# Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
|
|
if args.use_8bit_adam:
|
|
try:
|
|
import bitsandbytes as bnb
|
|
except ImportError:
|
|
raise ImportError("To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`.")
|
|
|
|
optimizer_class = bnb.optim.AdamW8bit
|
|
else:
|
|
optimizer_class = torch.optim.AdamW
|
|
|
|
# Optimizer creation
|
|
params_to_optimize = rdt.parameters()
|
|
optimizer = optimizer_class(
|
|
params_to_optimize,
|
|
lr=args.learning_rate,
|
|
betas=(args.adam_beta1, args.adam_beta2),
|
|
weight_decay=args.adam_weight_decay,
|
|
eps=args.adam_epsilon,
|
|
)
|
|
|
|
# Dataset and DataLoaders creation:
|
|
train_dataset = VLAConsumerDataset(
|
|
model_config_path=args.model_config_path, # TODO
|
|
config=config["dataset"],
|
|
tokenizer=tokenizer,
|
|
image_processor=image_processor,
|
|
num_cameras=config["common"]["num_cameras"],
|
|
img_history_size=config["common"]["img_history_size"],
|
|
dataset_type=args.dataset_type,
|
|
image_aug=args.image_aug,
|
|
cond_mask_prob=args.cond_mask_prob,
|
|
cam_ext_mask_prob=args.cam_ext_mask_prob,
|
|
state_noise_snr=args.state_noise_snr,
|
|
use_hdf5=args.load_from_hdf5,
|
|
use_precomp_lang_embed=args.precomp_lang_embed,
|
|
)
|
|
sample_dataset = VLAConsumerDataset(
|
|
model_config_path=args.model_config_path, # TODO
|
|
config=config["dataset"],
|
|
tokenizer=tokenizer,
|
|
image_processor=image_processor,
|
|
num_cameras=config["common"]["num_cameras"],
|
|
img_history_size=config["common"]["img_history_size"],
|
|
dataset_type=args.dataset_type,
|
|
image_aug=False,
|
|
cond_mask_prob=0,
|
|
cam_ext_mask_prob=-1,
|
|
state_noise_snr=None,
|
|
use_hdf5=args.load_from_hdf5,
|
|
use_precomp_lang_embed=args.precomp_lang_embed,
|
|
)
|
|
|
|
data_collator = DataCollatorForVLAConsumerDataset(tokenizer)
|
|
|
|
train_dataloader = torch.utils.data.DataLoader(
|
|
train_dataset,
|
|
batch_size=args.train_batch_size,
|
|
shuffle=True,
|
|
collate_fn=data_collator,
|
|
num_workers=args.dataloader_num_workers,
|
|
pin_memory=True,
|
|
persistent_workers=True,
|
|
)
|
|
sample_dataloader = torch.utils.data.DataLoader(
|
|
sample_dataset,
|
|
batch_size=args.sample_batch_size,
|
|
shuffle=True,
|
|
collate_fn=data_collator,
|
|
num_workers=args.dataloader_num_workers,
|
|
pin_memory=True,
|
|
persistent_workers=True,
|
|
)
|
|
|
|
# Scheduler and math around the number of training steps.
|
|
overrode_max_train_steps = False
|
|
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
|
|
if args.max_train_steps is None:
|
|
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
|
|
overrode_max_train_steps = True
|
|
|
|
lr_scheduler = get_scheduler(
|
|
args.lr_scheduler,
|
|
optimizer=optimizer,
|
|
num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
|
|
num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
|
|
num_cycles=args.lr_num_cycles,
|
|
power=args.lr_power,
|
|
)
|
|
|
|
# Prepare everything with our `accelerator`.
|
|
rdt, optimizer, train_dataloader, sample_dataloader, lr_scheduler = (accelerator.prepare(
|
|
rdt, optimizer, train_dataloader, sample_dataloader, lr_scheduler))
|
|
|
|
ema_rdt.to(accelerator.device, dtype=weight_dtype)
|
|
|
|
if text_encoder is not None:
|
|
text_encoder.to(accelerator.device, dtype=weight_dtype)
|
|
|
|
if vision_encoder is not None:
|
|
vision_encoder.vision_tower.to(accelerator.device, dtype=weight_dtype)
|
|
|
|
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
|
|
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
|
|
if overrode_max_train_steps:
|
|
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
|
|
# Afterwards we recalculate our number of training epochs
|
|
args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
|
|
|
|
# We need to initialize the trackers we use, and also store our configuration.
|
|
# The trackers initializes automatically on the main process.
|
|
if accelerator.is_main_process:
|
|
accelerator.init_trackers(
|
|
"VLA",
|
|
config=vars(args),
|
|
init_kwargs={"wandb": {
|
|
"name": f"RoboTwin_RDT_{args.CONFIG_NAME}",
|
|
}},
|
|
)
|
|
|
|
# Train!
|
|
total_batch_size = (args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps)
|
|
|
|
logger.info("***** Running training *****")
|
|
logger.info(f" Num examples = {len(train_dataset)}")
|
|
logger.info(f" Num batches each epoch = {len(train_dataloader)}")
|
|
logger.info(f" Num Epochs = {args.num_train_epochs}")
|
|
logger.info(f" Instantaneous batch size per device = {args.train_batch_size}")
|
|
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
|
|
logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")
|
|
logger.info(f" Total optimization steps = {args.max_train_steps}")
|
|
global_step = 0
|
|
first_epoch = 0
|
|
|
|
# Load from a pretrained checkpoint
|
|
if (args.resume_from_checkpoint is None and args.pretrained_model_name_or_path is not None
|
|
and os.path.isfile(args.pretrained_model_name_or_path)):
|
|
# Since EMA is deprecated, we do not load EMA from the pretrained checkpoint
|
|
logger.info("Loading from a pretrained checkpoint.")
|
|
checkpoint = torch.load(args.pretrained_model_name_or_path)
|
|
rdt.module.load_state_dict(checkpoint["module"])
|
|
|
|
# Potentially load in the weights and states from a previous save
|
|
if args.resume_from_checkpoint:
|
|
if args.resume_from_checkpoint != "latest":
|
|
path = os.path.basename(args.resume_from_checkpoint)
|
|
else:
|
|
# Get the mos recent checkpoint
|
|
dirs = os.listdir(args.output_dir)
|
|
dirs = [d for d in dirs if d.startswith("checkpoint")]
|
|
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
|
|
path = dirs[-1] if len(dirs) > 0 else None
|
|
|
|
if path is None:
|
|
accelerator.print(
|
|
f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.")
|
|
args.resume_from_checkpoint = None
|
|
else:
|
|
accelerator.print(f"Resuming from checkpoint {path}")
|
|
try:
|
|
accelerator.load_state(os.path.join(args.output_dir, path)) # load_module_strict=False
|
|
except:
|
|
# load deepspeed's state_dict
|
|
logger.info("Resuming training state failed. Attempting to only load from model checkpoint.")
|
|
checkpoint = torch.load(
|
|
os.path.join(
|
|
args.output_dir,
|
|
path,
|
|
"pytorch_model",
|
|
"mp_rank_00_model_states.pt",
|
|
))
|
|
rdt.module.load_state_dict(checkpoint["module"])
|
|
|
|
load_model(ema_rdt, os.path.join(args.output_dir, path, "ema", "model.safetensors"))
|
|
global_step = int(path.split("-")[1])
|
|
|
|
resume_global_step = global_step * args.gradient_accumulation_steps
|
|
first_epoch = global_step // num_update_steps_per_epoch
|
|
resume_step = resume_global_step % (num_update_steps_per_epoch * args.gradient_accumulation_steps)
|
|
|
|
# Only show the progress bar once on each machine.
|
|
progress_bar = tqdm(
|
|
range(global_step, args.max_train_steps),
|
|
disable=not accelerator.is_local_main_process,
|
|
)
|
|
progress_bar.set_description("Steps")
|
|
|
|
loss_for_log = {}
|
|
for epoch in range(first_epoch, args.num_train_epochs):
|
|
|
|
rdt.train()
|
|
|
|
# Set the progress_bar to correct position
|
|
if args.resume_from_checkpoint and epoch == first_epoch:
|
|
progress_bar.update(resume_step // args.gradient_accumulation_steps)
|
|
|
|
# Forward and backward...
|
|
for batch in train_dataloader:
|
|
with accelerator.accumulate(rdt):
|
|
images = batch["images"].to(dtype=weight_dtype)
|
|
states = batch["states"].to(dtype=weight_dtype) # (B, T, D_a)
|
|
# We only use the last state as input
|
|
states = states[:, -1:, :]
|
|
actions = batch["actions"].to(dtype=weight_dtype)
|
|
state_elem_mask = batch["state_elem_mask"].to(dtype=weight_dtype)
|
|
ctrl_freqs = batch["ctrl_freqs"]
|
|
|
|
with torch.no_grad():
|
|
batch_size, _, C, H, W = images.shape
|
|
image_embeds = vision_encoder(images.reshape(-1, C, H, W)).detach()
|
|
image_embeds = image_embeds.reshape((batch_size, -1, vision_encoder.hidden_size))
|
|
|
|
lang_attn_mask = batch["lang_attn_mask"]
|
|
text_embeds = (batch["lang_embeds"].to(
|
|
dtype=weight_dtype) if args.precomp_lang_embed else text_encoder(
|
|
input_ids=batch["input_ids"], attention_mask=lang_attn_mask)["last_hidden_state"].detach())
|
|
|
|
state_elem_mask = state_elem_mask.unsqueeze(1)
|
|
loss = rdt(
|
|
lang_tokens=text_embeds,
|
|
lang_attn_mask=lang_attn_mask,
|
|
img_tokens=image_embeds,
|
|
state_tokens=states,
|
|
action_gt=actions,
|
|
action_mask=state_elem_mask,
|
|
ctrl_freqs=ctrl_freqs,
|
|
)
|
|
|
|
accelerator.backward(loss)
|
|
if accelerator.sync_gradients:
|
|
params_to_clip = rdt.parameters()
|
|
accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
|
|
optimizer.step()
|
|
lr_scheduler.step()
|
|
optimizer.zero_grad(set_to_none=args.set_grads_to_none)
|
|
|
|
ema_model.step(accelerator.unwrap_model(rdt))
|
|
|
|
# Checks if the accelerator has performed an optimization step behind the scenes
|
|
if accelerator.sync_gradients:
|
|
progress_bar.update(1)
|
|
global_step += 1
|
|
|
|
if global_step % args.checkpointing_period == 0:
|
|
save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
|
|
accelerator.save_state(save_path)
|
|
ema_save_path = os.path.join(save_path, f"ema")
|
|
accelerator.save_model(ema_rdt, ema_save_path)
|
|
logger.info(f"Saved state to {save_path}")
|
|
|
|
if args.sample_period > 0 and global_step % args.sample_period == 0:
|
|
sample_loss_for_log = log_sample_res(
|
|
text_encoder,
|
|
vision_encoder,
|
|
rdt, # We do not use EMA currently
|
|
args,
|
|
accelerator,
|
|
weight_dtype,
|
|
sample_dataset.get_dataset_id2name(),
|
|
sample_dataloader,
|
|
logger,
|
|
)
|
|
logger.info(sample_loss_for_log)
|
|
accelerator.log(sample_loss_for_log, step=global_step)
|
|
|
|
logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
|
|
progress_bar.set_postfix(**logs)
|
|
logs.update(loss_for_log)
|
|
# logger.info(logs)
|
|
accelerator.log(logs, step=global_step)
|
|
|
|
if global_step >= args.max_train_steps:
|
|
break
|
|
|
|
# Create the pipeline using using the trained modules and save it.
|
|
accelerator.wait_for_everyone()
|
|
if accelerator.is_main_process:
|
|
accelerator.unwrap_model(rdt).save_pretrained(args.output_dir)
|
|
ema_save_path = os.path.join(args.output_dir, f"ema")
|
|
accelerator.save_model(ema_rdt, ema_save_path)
|
|
|
|
logger.info(f"Saved Model to {args.output_dir}")
|
|
|
|
if args.push_to_hub:
|
|
save_model_card(
|
|
repo_id,
|
|
base_model=args.pretrained_model_name_or_path,
|
|
repo_folder=args.output_dir,
|
|
)
|
|
upload_folder(
|
|
repo_id=repo_id,
|
|
folder_path=args.output_dir,
|
|
commit_message="End of training",
|
|
token=args.hub_token,
|
|
allow_patterns=["pytorch_model.bin", "*.json", "*.md"],
|
|
# ignore_patterns=["step_*", "epoch_*"],
|
|
)
|
|
|
|
accelerator.end_training()
|