Spaces:
Running on Zero
Running on Zero
File size: 4,953 Bytes
ccfee12 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import os, torch
from tqdm import tqdm
from accelerate import Accelerator
from .training_module import DiffusionTrainingModule
from src.training_module import MetaViewTrainingModule
from .logger import ModelLogger
from PIL import Image
def collate_fn(batch):
if len(batch) == 0:
return {}
collated = {}
keys = batch[0].keys()
for key in keys:
values = [sample[key] for sample in batch]
first_val = values[0]
if isinstance(first_val, torch.Tensor):
# 对于 Tensor，使用 torch.stack 沿第0维堆叠（要求所有张量形状相同）
collated[key] = torch.stack(values, dim=0)
elif isinstance(first_val, Image.Image):
collated[key] = values
elif isinstance(first_val, str):
collated[key] = values
else:
collated[key] = values
return collated
def launch_training_task(
accelerator: Accelerator,
dataset: torch.utils.data.Dataset,
model: MetaViewTrainingModule,
model_logger: ModelLogger,
learning_rate: float = 1e-5,
weight_decay: float = 1e-2,
num_workers: int = 1,
save_steps: int = None,
num_epochs: int = 1,
batch_size: int = 1,
args = None,
):
if args is not None:
learning_rate = args.learning_rate
weight_decay = args.weight_decay
num_workers = args.dataset_num_workers
save_steps = args.save_steps
num_epochs = args.num_epochs
optimizer = torch.optim.AdamW(model.trainable_modules(), lr=learning_rate, weight_decay=weight_decay)
scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)
# dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, collate_fn=lambda x: x[0], num_workers=num_workers)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn, num_workers=num_workers)
model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler)
def grad_hook(name):
def hook(grad):
# # print(f"{name} gradient norm:{grad.norm().item():.6f}")
# if grad.norm().item() > 1:
# print(f" gradient over 1: {name} {grad.norm().item():.6f}")
if torch.isnan(grad).any():
print(f"!!! NaN gradient: {name}")
if torch.isinf(grad).any():
print(f"!!! Inf gradient: {name}")
return grad
return hook
for name, param in model.named_parameters():
if param.requires_grad:
param.register_hook(grad_hook(name))
NaN_step = 0
for epoch_id in range(num_epochs):
for data in tqdm(dataloader):
with accelerator.accumulate(model):
# print(type(data))
# print(data["prompt"])
optimizer.zero_grad()
if dataset.load_from_cache:
loss = model({}, inputs=data)
else:
loss = model(data)
if torch.isnan(loss).any():
print(f"!!! Loss is NaN at step {model_logger.num_steps}! Skipping...")
NaN_step += 1
print(data["name"])
exit(0)
accelerator.backward(loss)
max_norm = 5.0
if accelerator.sync_gradients:
grad_norm = accelerator.clip_grad_norm_(model.parameters(), max_norm=max_norm)
if accelerator.is_main_process:
if grad_norm > 5.0:
print(f"gradient over 5: {grad_norm:.4f}")
optimizer.step()
model_logger.on_step_end(accelerator, model, save_steps)
scheduler.step()
if save_steps is None:
model_logger.on_epoch_end(accelerator, model, epoch_id)
model_logger.on_training_end(accelerator, model, save_steps)
def launch_data_process_task(
accelerator: Accelerator,
dataset: torch.utils.data.Dataset,
model: DiffusionTrainingModule,
model_logger: ModelLogger,
num_workers: int = 8,
args = None,
):
if args is not None:
num_workers = args.dataset_num_workers
dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, collate_fn=lambda x: x[0], num_workers=num_workers)
model, dataloader = accelerator.prepare(model, dataloader)
for data_id, data in enumerate(tqdm(dataloader)):
with accelerator.accumulate(model):
with torch.no_grad():
folder = os.path.join(model_logger.output_path, str(accelerator.process_index))
os.makedirs(folder, exist_ok=True)
save_path = os.path.join(model_logger.output_path, str(accelerator.process_index), f"{data_id}.pth")
data = model(data)
torch.save(data, save_path)
|