def set_model(self, inference=False):
self.model_type = 'linear' if 'phone' in self.task else 'rnn'
input_dim = int(self.config['downstream'][self.model_type]['input_dim']) if \
self.config['downstream'][self.model_type]['input_dim'] != 'None' else None
if 'mockingjay' in self.task:
self.mockingjay = Tester(self.mock_config, self.mock_paras)
if self.fine_tune and inference:
self.mockingjay.load = False # Do not load twice when testing the fine-tuned model, load only for fine-tune training
self.mockingjay.set_model(inference=True,
with_head=self.paras.with_head)
self.dr = self.mockingjay.dr
if input_dim is None:
input_dim = self.mock_config['mockingjay']['hidden_size']
elif 'apc' in self.task:
self.apc = get_apc_model(path=self.paras.apc_path)
if input_dim is None:
input_dim = self.mock_config['mockingjay'][
'hidden_size'] # use identical dim size for fair comparison
elif 'baseline' in self.task:
if input_dim is None:
input_dim = mel_dim
else:
raise NotImplementedError('Invalid Task!')
if self.model_type == 'linear':
self.classifier = LinearClassifier(
input_dim=input_dim,
class_num=self.dataloader.dataset.class_num,
task=self.task,
dconfig=self.config['downstream']['linear'],
sequencial=False).to(self.device)
elif self.model_type == 'rnn':
self.classifier = RnnClassifier(
input_dim=input_dim,
class_num=self.dataloader.dataset.class_num,
task=self.task,
dconfig=self.config['downstream']['rnn']).to(self.device)
if not inference and self.fine_tune:
# Setup Fine tune optimizer
self.mockingjay.mockingjay.train()
param_optimizer = list(
self.mockingjay.mockingjay.named_parameters()) + list(
self.classifier.named_parameters())
self.optimizer = get_mockingjay_optimizer(
params=param_optimizer,
lr=self.learning_rate,
warmup_proportion=self.config['optimizer']
['warmup_proportion'],
training_steps=self.total_steps)
elif not inference:
self.optimizer = Adam(self.classifier.parameters(),
lr=self.learning_rate,
betas=(0.9, 0.999))
self.classifier.train()
else:
self.classifier.eval()
if self.load: # This will be set to True by default when Tester is running set_model()
self.load_model(inference=inference)
def get_downstream_model(args, input_dim, class_num, config):
model_name = args.run.split('_')[-1].replace('frame', 'linear') # support names: ['linear', '1hidden', 'concat', 'utterance']
model_config = config['model'][model_name]
if args.task == 'speaker' and 'utterance' in args.run:
downstream_model = RnnClassifier(input_dim, class_num, model_config)
else:
downstream_model = LinearClassifier(input_dim, class_num, model_config)
return downstream_model
class Downstream_Solver(Solver):
''' Handler for complete training progress'''
def __init__(self, config, paras, task):
super(Downstream_Solver, self).__init__(config, paras)
# Downstream task the solver is solving
self.task = task
self.mock_paras = copy.deepcopy(paras)
self.mock_config = copy.deepcopy(config)
# path and directories
self.exp_name = self.exp_name.replace('mockingjay', task)
self.paras.ckpdir = paras.ckpdir.replace('mockingjay', task)
self.ckpdir = self.ckpdir.replace('mockingjay', task)
self.ckpt = os.path.join(paras.ckpdir, paras.dckpt)
# modify log directory
paras.logdir = paras.logdir.replace('mockingjay', task)
# model
self.model_type = config['downstream']['model_type']
self.load_model_list = config['downstream']['load_model_list']
self.fine_tune = paras.fine_tune
self.run_mockingjay = paras.run_mockingjay
self.run_apc = paras.run_apc
if self.fine_tune:
assert (
self.run_mockingjay
), 'Use `--run_mockingjay` to fine-tune the mockingjay model.'
assert (not self.run_apc
), 'Fine tuning only supports the mockingjay model.'
assert (
not self.paras.with_head
), 'Fine tuning only supports the mockingjay model, not with head.'
assert (not (self.run_mockingjay and self.run_apc)
), 'Mockingjay and Apc can not run at the same time!'
if self.run_mockingjay and self.paras.with_head:
self.verbose('Using Mockingjay representations from head.')
elif self.run_mockingjay and self.fine_tune:
self.verbose('Fine-tuning on Mockingjay representations.')
elif self.run_mockingjay:
self.verbose('Using Mockingjay representations.')
def load_data(self, split='train', load='phone'):
''' Load date for training / testing'''
assert (load in [
'phone', 'cpc_phone', 'sentiment', 'speaker', 'speaker_large'
]), 'Unsupported dataloader!'
if load == 'phone' or load == 'cpc_phone' or load == 'speaker_large':
if split == 'train':
self.verbose('Loading source data from ' +
str(self.config['dataloader']['train_set']) +
' from ' + self.config['dataloader']['data_path'])
if load == 'phone' or load == 'cpc_phone':
self.verbose('Loading phone data from ' +
str(self.config['dataloader']['train_set']) +
' from ' +
self.config['dataloader']['phone_path'])
elif split == 'test':
if load != 'cpc_phone':
self.verbose('Loading testing data ' +
str(self.config['dataloader']['test_set']) +
' from ' +
self.config['dataloader']['data_path'])
if load == 'phone':
self.verbose('Loading label data ' +
str(self.config['dataloader']['test_set']) +
' from ' +
self.config['dataloader']['phone_path'])
elif load == 'cpc_phone':
self.verbose('Loading label data from ' +
self.config['dataloader']['phone_path'])
else:
raise NotImplementedError('Invalid `split` argument!')
elif load == 'speaker':
if split == 'train':
self.verbose('Loading source data from ' +
str(self.config['dataloader']
['train_set']).replace('360', '100') +
' from ' + self.config['dataloader']['data_path'])
elif split == 'test':
self.verbose('Loading testing data ' +
str(self.config['dataloader']
['test_set']).replace('360', '100') +
' from ' + self.config['dataloader']['data_path'])
else:
raise NotImplementedError('Invalid `split` argument!')
elif load == 'sentiment':
target = self.config['dataloader']['sentiment_config']['dataset']
sentiment_path = self.config['dataloader']['sentiment_config'][
target]['path']
self.verbose(f'Loading {split} data from {sentiment_path}')
else:
raise NotImplementedError('Unsupported downstream tasks.')
setattr(self, 'dataloader', get_Dataloader(split, load=load, use_gpu=self.paras.gpu, \
run_mockingjay=self.run_mockingjay, mock_config=self.config['mockingjay'], \
**self.config['dataloader']))
def set_model(self, inference=False):
input_dim = int(self.config['downstream'][self.model_type]['input_dim']) if \
self.config['downstream'][self.model_type]['input_dim'] != 'None' else None
if 'mockingjay' in self.task:
self.mockingjay = Tester(self.mock_config, self.mock_paras)
if self.fine_tune and inference:
self.mockingjay.load = False # Do not load twice when testing the fine-tuned model, load only for fine-tune training
self.mockingjay.set_model(inference=True,
with_head=self.paras.with_head)
self.dr = self.mockingjay.dr
if input_dim is None:
input_dim = self.mock_config['mockingjay']['hidden_size']
elif 'apc' in self.task:
self.apc = get_apc_model(path=self.paras.apc_path)
if input_dim is None:
input_dim = self.mock_config['mockingjay'][
'hidden_size'] # use identical dim size for fair comparison
elif 'baseline' in self.task:
if input_dim is None:
if 'input_dim' in self.mock_config['mockingjay']:
input_dim = self.mock_config['mockingjay']['input_dim']
self.verbose(
'Using `input_dim` setting from config for downstream model.'
)
else:
input_dim = fmllr_dim if 'fmllr' in self.config[
'dataloader'][
'data_path'] else mfcc_dim if 'mfcc' in self.config[
'dataloader']['data_path'] else mel_dim
else:
raise NotImplementedError('Invalid Task!')
if self.model_type == 'linear':
self.classifier = LinearClassifier(
input_dim=input_dim,
class_num=self.dataloader.dataset.class_num,
task=self.task,
dconfig=self.config['downstream']['linear']).to(self.device)
elif self.model_type == 'rnn':
self.classifier = RnnClassifier(
input_dim=input_dim,
class_num=self.dataloader.dataset.class_num,
task=self.task,
dconfig=self.config['downstream']['rnn']).to(self.device)
if not inference and self.fine_tune:
# Setup Fine tune optimizer
self.mockingjay.mockingjay.train()
param_optimizer = list(
self.mockingjay.mockingjay.named_parameters()) + list(
self.classifier.named_parameters())
self.optimizer = get_mockingjay_optimizer(
params=param_optimizer,
lr=self.learning_rate,
warmup_proportion=self.config['optimizer']
['warmup_proportion'],
training_steps=self.total_steps)
elif not inference:
self.optimizer = Adam(self.classifier.parameters(),
lr=self.learning_rate,
betas=(0.9, 0.999))
self.classifier.train()
else:
self.classifier.eval()
if self.load: # This will be set to True by default when Tester is running set_model()
self.load_model(inference=inference)
def save_model(self, name, model_all=True, assign_name=None):
if model_all:
all_states = {
'Classifier':
self.classifier.state_dict(),
'Mockingjay':
self.mockingjay.mockingjay.state_dict()
if self.fine_tune else None,
'Optimizer':
self.optimizer.state_dict(),
'Global_step':
self.global_step,
'Settings': {
'Config': self.config,
'Paras': self.paras,
},
}
else:
all_states = {
'Classifier': self.classifier.state_dict(),
'Settings': {
'Config': self.config,
'Paras': self.paras,
},
}
if assign_name is not None:
model_path = f'{self.ckpdir}/{assign_name}.ckpt'
torch.save(all_states, model_path)
return
new_model_path = '{}/{}-{}.ckpt'.format(self.ckpdir, name,
self.global_step)
torch.save(all_states, new_model_path)
self.model_kept.append(new_model_path)
if len(self.model_kept) >= self.max_keep:
os.remove(self.model_kept[0])
self.model_kept.pop(0)
def load_model(self, inference=False):
self.verbose('Load model from {}'.format(self.ckpt))
all_states = torch.load(self.ckpt, map_location='cpu')
if 'Classifier' in self.load_model_list:
try:
self.classifier.load_state_dict(all_states['Classifier'])
self.verbose('[Classifier] - Loaded')
except:
self.verbose('[Classifier - X]')
if 'Optimizer' in self.load_model_list and not inference:
try:
self.optimizer.load_state_dict(all_states['Optimizer'])
for state in self.optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.cuda()
self.verbose('[Optimizer] - Loaded')
except:
self.verbose('[Optimizer - X]')
if 'Global_step' in self.load_model_list:
try:
self.global_step = all_states['Global_step']
self.verbose('[Global_step] - Loaded')
except:
self.verbose('[Global_step - X]')
if self.fine_tune:
try:
self.verbose(
'@ Downstream, [Fine-Tuned Mockingjay] - Loading from Upstream Tester...'
)
self.mockingjay.load_model(inference=inference,
from_path=self.ckpt)
self.verbose('@ Downstream, [Fine-Tuned Mockingjay] - Loaded')
except:
self.verbose('[Fine-Tuned Mockingjay] - X')
self.verbose('Model loading complete!')
def set_model(self, inference=False):
input_dim = int(self.config['downstream'][self.model_type]['input_dim']) if \
self.config['downstream'][self.model_type]['input_dim'] != 'None' else None
if 'transformer' in self.task:
self.upstream_tester = Tester(self.upstream_config,
self.upstream_paras)
if self.fine_tune and inference:
self.upstream_tester.load = False # During inference on fine-tuned model, load with `load_downstream_model()`
self.upstream_tester.set_model(
inference=True, with_head=self.paras.with_head
) # inference should be set True so upstream solver won't create optimizer
self.dr = self.upstream_tester.dr
if input_dim is None:
input_dim = self.transformer_config['hidden_size']
elif 'apc' in self.task:
self.apc = get_apc_model(path=self.paras.apc_path)
if input_dim is None:
input_dim = self.transformer_config[
'hidden_size'] # use identical dim size for fair comparison
elif 'baseline' in self.task:
if input_dim is None:
if 'input_dim' in self.transformer_config:
input_dim = self.transformer_config['input_dim']
else:
raise ValueError(
'Please update your config file to include the attribute `input_dim`.'
)
else:
raise NotImplementedError('Invalid Task!')
if self.model_type == 'linear':
self.classifier = LinearClassifier(
input_dim=input_dim,
class_num=self.dataloader.dataset.class_num,
dconfig=self.config['downstream']['linear']).to(self.device)
elif self.model_type == 'rnn':
self.classifier = RnnClassifier(
input_dim=input_dim,
class_num=self.dataloader.dataset.class_num,
dconfig=self.config['downstream']['rnn']).to(self.device)
if not inference and self.fine_tune:
# Setup Fine tune optimizer
self.upstream_tester.transformer.train()
param_optimizer = list(
self.upstream_tester.transformer.named_parameters()) + list(
self.classifier.named_parameters())
self.optimizer = get_optimizer(
params=param_optimizer,
lr=self.learning_rate,
warmup_proportion=self.config['optimizer']
['warmup_proportion'],
training_steps=self.total_steps)
elif not inference:
self.optimizer = Adam(self.classifier.parameters(),
lr=self.learning_rate,
betas=(0.9, 0.999))
self.classifier.train()
else:
self.classifier.eval()
if self.load: # This will be set to True by default when Tester is running set_model()
self.load_downstream_model(inference=inference)