eden-emu/opus
14
0
Fork 0
mirror of https://github.com/xiph/opus.git synced 2025-05-15 16:08:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Adds --chunks-per-offset option to train_rdovae.py

Browse source
This commit is contained in:
Jean-Marc Valin 2025-03-28 19:28:14 -04:00
parent 1ca6933ac4
commit af6dbd84a8
No known key found for this signature in database
GPG key ID: 5E5DD9A36F9189C8
2 changed files with 6 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

4
dnn/torch/rdovae/rdovae/rdovae.py
View file

@ -551,6 +551,7 @@ class RDOVAE(nn.Module):
cond_size2, cond_size2,
state_dim=24, state_dim=24,
split_mode='split', split_mode='split',
chunks_per_offset=4,
clip_weights=False, clip_weights=False,
pvq_num_pulses=82, pvq_num_pulses=82,
state_dropout_rate=0, state_dropout_rate=0,
@ -564,6 +565,7 @@ class RDOVAE(nn.Module):
self.cond_size = cond_size self.cond_size = cond_size
self.cond_size2 = cond_size2 self.cond_size2 = cond_size2
self.split_mode = split_mode self.split_mode = split_mode
self.chunks_per_offset = chunks_per_offset
self.state_dim = state_dim self.state_dim = state_dim
self.pvq_num_pulses = pvq_num_pulses self.pvq_num_pulses = pvq_num_pulses
self.state_dropout_rate = state_dropout_rate self.state_dropout_rate = state_dropout_rate
@ -670,7 +672,7 @@ class RDOVAE(nn.Module):
states_q = states_q * mask states_q = states_q * mask
# decoder # decoder
chunks = self.get_decoder_chunks(z.size(1), mode=self.split_mode) chunks = self.get_decoder_chunks(z.size(1), mode=self.split_mode, chunks_per_offset=self.chunks_per_offset)
outputs_hq = [] outputs_hq = []
outputs_sq = [] outputs_sq = []

5
dnn/torch/rdovae/train_rdovae.py
View file

@ -60,7 +60,8 @@ training_group = parser.add_argument_group(title="training parameters")
training_group.add_argument('--batch-size', type=int, help="batch size, default: 32", default=32) training_group.add_argument('--batch-size', type=int, help="batch size, default: 32", default=32)
training_group.add_argument('--lr', type=float, help='learning rate, default: 3e-4', default=3e-4) training_group.add_argument('--lr', type=float, help='learning rate, default: 3e-4', default=3e-4)
training_group.add_argument('--epochs', type=int, help='number of training epochs, default: 100', default=100) training_group.add_argument('--epochs', type=int, help='number of training epochs, default: 100', default=100)
training_group.add_argument('--sequence-length', type=int, help='sequence length, needs to be divisible by 4, default: 256', default=256) training_group.add_argument('--sequence-length', type=int, help='sequence length, needs to be divisible by chunks_per_offset, default: 400', default=400)
training_group.add_argument('--chunks-per-offset', type=int, help='chunks per offset', default=4)
training_group.add_argument('--lr-decay-factor', type=float, help='learning rate decay factor, default: 2.5e-5', default=2.5e-5) training_group.add_argument('--lr-decay-factor', type=float, help='learning rate decay factor, default: 2.5e-5', default=2.5e-5)
training_group.add_argument('--split-mode', type=str, choices=['split', 'random_split'], help='splitting mode for decoder input, default: split', default='split') training_group.add_argument('--split-mode', type=str, choices=['split', 'random_split'], help='splitting mode for decoder input, default: split', default='split')
training_group.add_argument('--enable-first-frame-loss', action='store_true', default=False, help='enables dedicated distortion loss on first 4 decoder frames') training_group.add_argument('--enable-first-frame-loss', action='store_true', default=False, help='enables dedicated distortion loss on first 4 decoder frames')
@ -120,7 +121,7 @@ feature_file = args.features
# model # model
checkpoint['model_args'] = (num_features, latent_dim, quant_levels, cond_size, cond_size2) checkpoint['model_args'] = (num_features, latent_dim, quant_levels, cond_size, cond_size2)
checkpoint['model_kwargs'] = {'state_dim': state_dim, 'split_mode' : split_mode, 'pvq_num_pulses': args.pvq_num_pulses, 'state_dropout_rate': args.state_dropout_rate, 'softquant': softquant} checkpoint['model_kwargs'] = {'state_dim': state_dim, 'split_mode' : split_mode, 'pvq_num_pulses': args.pvq_num_pulses, 'state_dropout_rate': args.state_dropout_rate, 'softquant': softquant, 'chunks_per_offset': args.chunks_per_offset}
model = RDOVAE(*checkpoint['model_args'], **checkpoint['model_kwargs']) model = RDOVAE(*checkpoint['model_args'], **checkpoint['model_kwargs'])
if type(args.initial_checkpoint) != type(None): if type(args.initial_checkpoint) != type(None):