77 lines
2.9 KiB
Python
77 lines
2.9 KiB
Python
# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
|
|
# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
|
|
|
|
# Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
|
|
# The below software and/or models in this distribution may have been
|
|
# modified by THL A29 Limited ("Tencent Modifications").
|
|
# All Tencent Modifications are Copyright (C) THL A29 Limited.
|
|
|
|
# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
|
|
# except for the third-party components listed below.
|
|
# Hunyuan 3D does not impose any additional limitations beyond what is outlined
|
|
# in the repsective licenses of these third-party components.
|
|
# Users must comply with all terms and conditions of original licenses of these third-party
|
|
# components and must ensure that the usage of the third party components adheres to
|
|
# all relevant laws and regulations.
|
|
|
|
# For avoidance of doubts, Hunyuan 3D means the large language models and
|
|
# their software and algorithms, including trained model weights, parameters (including
|
|
# optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
|
|
# fine-tuning enabling code and other elements of the foregoing made publicly available
|
|
# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
|
|
|
|
import os
|
|
import time
|
|
import random
|
|
import numpy as np
|
|
import torch
|
|
from functools import wraps
|
|
|
|
def seed_everything(seed):
|
|
'''
|
|
seed everthing
|
|
'''
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
os.environ["PL_GLOBAL_SEED"] = str(seed)
|
|
|
|
def timing_decorator(category: str):
|
|
'''
|
|
timing_decorator: record time
|
|
'''
|
|
def decorator(func):
|
|
func.call_count = 0
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
start_time = time.time()
|
|
result = func(*args, **kwargs)
|
|
end_time = time.time()
|
|
elapsed_time = end_time - start_time
|
|
func.call_count += 1
|
|
print(f"[HunYuan3D]-[{category}], cost time: {elapsed_time:.4f}s") # huiwen
|
|
return result
|
|
return wrapper
|
|
return decorator
|
|
|
|
def auto_amp_inference(func):
|
|
'''
|
|
with torch.cuda.amp.autocast()" # deprecated.
|
|
xxx
|
|
'''
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
with torch.amp.autocast('cuda'):
|
|
output = func(*args, **kwargs)
|
|
return output
|
|
return wrapper
|
|
|
|
def get_parameter_number(model):
|
|
total_num = sum(p.numel() for p in model.parameters())
|
|
trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
|
return {'Total': total_num, 'Trainable': trainable_num}
|
|
|
|
def set_parameter_grad_false(model):
|
|
for p in model.parameters():
|
|
p.requires_grad = False
|