nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-05-21 11:59:01 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

[fix] release av threads

Browse source
This commit is contained in:
zhengzangw 2024-05-30 08:49:39 +00:00
parent 687f603810
commit ee1354a7e0
2 changed files with 34 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
opensora/datasets/read_video.py
View file

@ -1,3 +1,4 @@
import gc
import math
import os
from fractions import Fraction
@ -70,38 +71,42 @@ def read_video(
audio_frames = []
audio_timebase = _video_opt.default_timebase
container = av.open(filename, metadata_errors="ignore")
try:
with av.open(filename, metadata_errors="ignore") as container:
if container.streams.audio:
audio_timebase = container.streams.audio[0].time_base
if container.streams.video:
video_frames = _read_from_stream(
container,
start_pts,
end_pts,
pts_unit,
container.streams.video[0],
{"video": 0},
)
video_fps = container.streams.video[0].average_rate
# guard against potentially corrupted files
if video_fps is not None:
info["video_fps"] = float(video_fps)
if container.streams.audio:
audio_frames = _read_from_stream(
container,
start_pts,
end_pts,
pts_unit,
container.streams.audio[0],
{"audio": 0},
)
info["audio_fps"] = container.streams.audio[0].rate
if container.streams.audio:
audio_timebase = container.streams.audio[0].time_base
if container.streams.video:
video_frames = _read_from_stream(
container,
start_pts,
end_pts,
pts_unit,
container.streams.video[0],
{"video": 0},
)
video_fps = container.streams.video[0].average_rate
# guard against potentially corrupted files
if video_fps is not None:
info["video_fps"] = float(video_fps)
if container.streams.audio:
audio_frames = _read_from_stream(
container,
start_pts,
end_pts,
pts_unit,
container.streams.audio[0],
{"audio": 0},
)
info["audio_fps"] = container.streams.audio[0].rate
except av.AVError:
# TODO raise a warning?
pass
finally:
container.close()
del container
# NOTE: manually garbage collect to close pyav threads
gc.collect()
vframes_list = [frame.to_rgb().to_ndarray() for frame in video_frames]
aframes_list = [frame.to_ndarray() for frame in audio_frames]

6
scripts/train.py
View file

@ -1,3 +1,4 @@
import gc
import os
from copy import deepcopy
from datetime import timedelta
@ -19,6 +20,7 @@ from opensora.registry import DATASETS, MODELS, SCHEDULERS, build_module
from opensora.utils.ckpt_utils import load, model_gathering, model_sharding, record_model_param_shape, save
from opensora.utils.config_utils import define_experiment_workspace, parse_configs, save_training_config
from opensora.utils.misc import (
Timer,
all_reduce_mean,
create_logger,
create_tensorboard_writer,
@ -26,7 +28,6 @@ from opensora.utils.misc import (
get_model_numel,
requires_grad,
to_torch_dtype,
Timer
)
from opensora.utils.train_utils import MaskGenerator, create_colossalai_plugin, update_ema
@ -104,8 +105,6 @@ def main():
**dataloader_args,
)
num_steps_per_epoch = len(dataloader)
dataiter = iter(dataloader)
next(dataiter)
# ======================================================
# 3. build model
@ -365,7 +364,6 @@ def main():
save_dir,
)
log_str = f"Rank {dist.get_rank()} | Epoch {epoch} | Step {step} | "
for timer in timer_list:
log_str += f"{timer.name}: {timer.elapsed_time:.3f}s | "