def __init__(self,
model,
loss_function,
train_data,
valid_data,
dicts,
opt,
setup_optimizer=True):
super().__init__(model, loss_function, train_data, valid_data, dicts,
opt)
if opt.lfv_multilingual or opt.lid_loss:
from onmt.models.speech_recognizer.lid_loss import CrossEntropyLIDLoss
lid_loss = CrossEntropyLIDLoss(opt.n_languages,
opt.label_smoothing,
opt.fast_xentropy)
self.loss_function.add_loss_function(lid_loss, 'lid_loss')
self.n_gpus = len(self.opt.gpus)
if self.cuda:
torch.cuda.set_device(self.opt.gpus[0])
if self.opt.seed >= 0:
torch.manual_seed(self.opt.seed)
self.loss_function = self.loss_function.cuda()
self.model = self.model.cuda()
if setup_optimizer:
self.optim = onmt.Optim(opt)
self.optim.set_parameters(self.model.parameters())
if not self.opt.fp16:
opt_level = "O0"
keep_batchnorm_fp32 = False
elif self.opt.fp16_mixed:
opt_level = "O1"
keep_batchnorm_fp32 = None
else:
opt_level = "O2"
keep_batchnorm_fp32 = False
if self.cuda:
# print(234)
self.model, self.optim.optimizer = amp.initialize(
self.model,
self.optim.optimizer,
opt_level=opt_level,
keep_batchnorm_fp32=keep_batchnorm_fp32,
loss_scale="dynamic",
verbosity=1 if self.opt.verbose else 0)
# An ugly hack to switch between align right and align left
if hasattr(self.model, 'relative'):
if self.model.relative:
self.train_data.src_align_right = True
self.train_data.tgt_align_right = False
self.valid_data.src_align_right = True
self.valid_data.tgt_align_right = False
self.valid_data.tgt_align_right = False
def __init__(self, model, loss_function, train_data, valid_data, dicts,
opt):
# self.model = model
init_model_parameters(model, opt)
self.model = model
if opt.load_encoder_from:
self.load_encoder_weight(opt.load_encoder_from)
if opt.load_decoder_from:
self.load_decoder_weight(opt.load_decoder_from)
self.model = None
self.train_data = train_data
self.valid_data = valid_data
self.dicts = dicts
self.opt = opt
self.cuda = (len(opt.gpus) >= 1 and opt.gpus[0] >= 0)
self.loss_function = loss_function
self.start_time = 0
# setting up models and others
if opt.lfv_multilingual:
from onmt.models.speech_recognizer.lid_loss import CrossEntropyLIDLoss
lid_loss = CrossEntropyLIDLoss(opt.n_languages,
opt.label_smoothing,
opt.fast_xentropy)
self.loss_function.add_loss_function(lid_loss, 'lid_loss')
self.model_wrapper = MultiprocessingRunner(opt,
model,
loss_function,
device_ids=opt.gpus)
def __init__(self, device, train_data, valid_data, dicts, opt, setup_optimizer=True):
"""
:param model:
:param device: int (GPU id)
:param loss_function:
:param train_data:
:param valid_data:
:param dicts:
:param opt:
"""
self.device = device
opt.node_rank = 0
opt.nodes = 1
self.world_size = len(opt.gpus)
# in the case of single node distributed, it should equal self.device
self.rank = self.device
# make a group to later use with self.all_reduce
self.group = dist.group.WORLD
self.print("[INFO] Training Options:", opt)
if self.world_size > 1:
dist.init_process_group(backend='nccl', init_method='env://', world_size=self.world_size, rank=self.rank)
self.model = None
if self.rank == 0:
self.train_data = train_data
self.valid_data = valid_data
else:
# Do we really need to deepcopy the data instances (which could cause memory leak easily)
self.train_data = copy.deepcopy(train_data)
self.valid_data = copy.deepcopy(valid_data)
self.dicts = dicts
self.opt = opt
self.cuda = (len(opt.gpus) >= 1 and opt.gpus[0] >= 0)
assert self.cuda, "[ERROR] Training is only available on GPUs."
self.start_time = 0
# setting up models and others
if opt.lfv_multilingual:
from onmt.models.speech_recognizer.lid_loss import CrossEntropyLIDLoss
lid_loss = CrossEntropyLIDLoss(opt.n_languages, opt.label_smoothing, opt.fast_xentropy)
self.loss_function.add_loss_function(lid_loss, 'lid_loss')
torch.manual_seed(self.opt.seed)
# note: we must start creating models after ccreating the processes
# for some reason passing a pre-created model to a process creates a "pickle" error
if not opt.fusion:
if self.is_main():
print("[INFO] Building models .... ", flush=True)
model = build_model(opt, dicts)
""" Building the loss function """
if opt.ctc_loss > 0.0:
from onmt.speech.ctc_loss import CTC
self.ctc_loss_function = CTC(0.0, reduce=True)
if opt.nce:
from onmt.modules.nce.nce_loss import NCELoss
loss_function = NCELoss(opt.model_size, dicts['tgt'].size(), noise_ratio=opt.nce_noise,
logz=9, label_smoothing=opt.label_smoothing)
else:
loss_function = NMTLossFunc(opt.model_size, dicts['tgt'].size(),
label_smoothing=opt.label_smoothing,
mirror=opt.mirror_loss,
fast_xentropy=opt.fast_xentropy)
# This function replaces modules with the more optimized counterparts so that it can run faster
# Currently exp with LayerNorm
if not opt.memory_profiling:
# distributed is required to convert BatchNorm to SyncBatchNorm for DDP
optimize_model(model, distributed=(self.world_size > 1))
init_model_parameters(model, opt)
self.model = model
self.loss_function = loss_function
self.grad_scaler = torch.cuda.amp.GradScaler()
if opt.load_from:
checkpoint = torch.load(opt.load_from, map_location=lambda storage, loc: storage)
self.model.load_state_dict(checkpoint['model'])
if 'scaler' in checkpoint and checkpoint['scaler'] is not None:
self.grad_scaler.load_state_dict(checkpoint['scaler'])
if self.cuda:
torch.cuda.set_device(self.device)
self.loss_function = self.loss_function.cuda(device=self.device)
self.model = self.model.cuda(device=self.device)
if opt.ctc_loss > 0.0:
self.ctc_loss_function = self.ctc_loss_function.cuda(device=self.device)
# Ensure that the distributed copies have the same initial parameters
# Manual seed may not work the same for different GPU models.
# if self.world_size > 1:
# params = [p for p in self.model.parameters()]
#
# with torch.no_grad():
# if not self.is_main():
# # zero everything except for the main model
# for p in params:
# p.zero_()
# else:
# for p in params:
# p.add_(0)
#
# # run all_reduce to ensure that all models have exactly the same parameters
# if self.world_size > 1:
# params = [p for p in self.model.parameters()]
# all_reduce_and_rescale_tensors(params, 1)
if setup_optimizer:
self.optim = onmt.Optim(opt)
self.optim.set_parameters(self.model.parameters())
if self.is_main():
print("[INFO] Optimizer: ", self.optim.optimizer)
if opt.load_from:
if 'optim' in checkpoint and checkpoint['optim'] is not None and not opt.reset_optim:
self.optim.load_state_dict(checkpoint['optim'])
if self.world_size > 1:
# find_unused_parameters may be required for dropped layer (parameters that are not connected to
# any particular graph)
find_unused_parameters = True
self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[self.rank],
output_device=self.rank,
find_unused_parameters=find_unused_parameters)
print("[INFO] Process %d ready." % self.rank, flush=True)
def __init__(self,
device,
train_data,
valid_data,
dicts,
opt,
setup_optimizer=True):
"""
:param model:
:param device: int (GPU id)
:param loss_function:
:param train_data:
:param valid_data:
:param dicts:
:param opt:
"""
# self.model = model
# self.model = model
# self.loss_function = loss_function
self.device = device
opt.node_rank = 0
opt.nodes = 1
self.world_size = len(opt.gpus)
# in the case of single node distributed, it should equal self.device
self.rank = self.device
# make a group to later use with dist.all_reduce
self.group = dist.group.WORLD
self.print("[INFO] Training Options:", opt)
dist.init_process_group(backend='nccl',
init_method='env://',
world_size=self.world_size,
rank=self.rank)
self.model = None
if self.rank == 0:
self.train_data = train_data
self.valid_data = valid_data
else:
self.train_data = copy.deepcopy(train_data)
self.valid_data = copy.deepcopy(valid_data)
self.dicts = dicts
self.opt = opt
self.cuda = (len(opt.gpus) >= 1 and opt.gpus[0] >= 0)
assert self.cuda, "[ERROR] Training is only available on GPUs."
self.start_time = 0
# setting up models and others
if opt.lfv_multilingual:
from onmt.models.speech_recognizer.lid_loss import CrossEntropyLIDLoss
lid_loss = CrossEntropyLIDLoss(opt.n_languages,
opt.label_smoothing,
opt.fast_xentropy)
self.loss_function.add_loss_function(lid_loss, 'lid_loss')
torch.manual_seed(self.opt.seed)
# note: we must start creating models after ccreating the processes
# for some reason passing a pre-created model to a process creates a "pickle" error
if not opt.fusion:
if self.is_main():
print("BUILDING MODEL .... ", flush=True)
model = build_model(opt, dicts)
""" Building the loss function """
if opt.ctc_loss != 0:
loss_function = NMTAndCTCLossFunc(
dicts['tgt'].size(),
label_smoothing=opt.label_smoothing,
ctc_weight=opt.ctc_loss)
elif opt.nce:
from onmt.modules.nce.nce_loss import NCELoss
loss_function = NCELoss(opt.model_size,
dicts['tgt'].size(),
noise_ratio=opt.nce_noise,
logz=9,
label_smoothing=opt.label_smoothing)
else:
loss_function = NMTLossFunc(
opt.model_size,
dicts['tgt'].size(),
label_smoothing=opt.label_smoothing,
mirror=opt.mirror_loss,
fast_xentropy=opt.fast_xentropy)
# This function replaces modules with the more optimized counterparts so that it can run faster
# Currently exp with LayerNorm
if not opt.memory_profiling:
# distributed is required to convert BatchNorm to SyncBatchNorm for DDP
optimize_model(model, distributed=(self.world_size > 1))
# optimize_model(model)
init_model_parameters(model, opt)
self.model = model
self.loss_function = loss_function
if self.cuda:
torch.cuda.set_device(self.device)
self.loss_function = self.loss_function.cuda(device=self.device)
self.model = self.model.cuda(device=self.device)
# Ensure that the distributed copies have the same initial parameters
# Manual seed may not work the same for different GPU models.
if self.world_size > 1:
params = [p for p in self.model.parameters()]
with torch.no_grad():
if not self.is_main():
for p in params:
p.zero_()
else:
for p in params:
p.add_(0)
if self.world_size > 1:
params = [p for p in self.model.parameters()]
all_reduce_and_rescale_tensors(params, 1)
if setup_optimizer:
self.optim = onmt.Optim(opt)
self.optim.set_parameters(self.model.parameters())
if self.is_main():
print("[INFO] Optimizer: ", self.optim.optimizer)
if not self.opt.fp16:
opt_level = "O0"
keep_batchnorm_fp32 = False
elif self.opt.fp16_mixed:
opt_level = "O1"
keep_batchnorm_fp32 = None
else:
opt_level = "O2"
keep_batchnorm_fp32 = False
if self.cuda:
self.model, self.optim.optimizer = amp.initialize(
self.model,
self.optim.optimizer,
opt_level=opt_level,
keep_batchnorm_fp32=keep_batchnorm_fp32,
loss_scale="dynamic",
verbosity=1 if self.opt.verbose else 0)
# wrap the model into DDP after initializing by amp
if self.world_size > 1:
"""
delay_allreduce is required to avoid allreduce error during backward pass
"""
self.model = DDP(self.model,
delay_allreduce=True,
gradient_average=False)
# torch DDP is more likely to work with the official amp autocast
# self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[self.rank],
# output_device=self.rank,
# find_unused_parameters=True)
print("[INFO] Process %d ready." % self.rank, flush=True)
def __init__(self,
device,
train_data,
valid_data,
dicts,
opt,
setup_optimizer=True):
"""
:param model:
:param device: int (GPU id)
:param loss_function:
:param train_data:
:param valid_data:
:param dicts:
:param opt:
"""
self.device = device
opt.node_rank = 0
opt.nodes = 1
self.world_size = len(opt.gpus)
# in the case of single node distributed, it should equal self.device
self.rank = self.device
# make a group to later use with self.all_reduce
self.group = dist.group.WORLD
self.print("[INFO] Training Options:", opt)
if self.world_size > 1:
dist.init_process_group(backend='nccl',
init_method='env://',
world_size=self.world_size,
rank=self.rank)
self.model = None
if self.rank == 0:
self.train_data = train_data
self.valid_data = valid_data
else:
# Do we really need to deepcopy the data instances (which could cause memory leak easily)
self.train_data = copy.deepcopy(train_data)
self.valid_data = copy.deepcopy(valid_data)
self.dicts = dicts
self.opt = opt
self.cuda = (len(opt.gpus) >= 1 and opt.gpus[0] >= 0)
assert self.cuda, "[ERROR] Training is only available on GPUs."
self.start_time = 0
# setting up models and others
if opt.lfv_multilingual:
from onmt.models.speech_recognizer.lid_loss import CrossEntropyLIDLoss
lid_loss = CrossEntropyLIDLoss(opt.n_languages,
opt.label_smoothing,
opt.fast_xentropy)
self.loss_function.add_loss_function(lid_loss, 'lid_loss')
torch.manual_seed(self.opt.seed)
# note: we must start creating models after ccreating the processes
# for some reason passing a pre-created model to a process creates a "pickle" error
if not opt.fusion:
if self.is_main():
print("[INFO] Building models .... ", flush=True)
model = build_model(opt, dicts)
""" Building the loss function """
if opt.ctc_loss > 0.0:
from onmt.speech.ctc_loss import CTC
self.ctc_loss_function = CTC(0.0, reduce=True)
if opt.nce:
from onmt.modules.nce.nce_loss import NCELoss
loss_function = NCELoss(opt.model_size,
dicts['tgt'].size(),
noise_ratio=opt.nce_noise,
logz=9,
label_smoothing=opt.label_smoothing)
else:
loss_function = NMTLossFunc(
opt.model_size,
dicts['tgt'].size(),
label_smoothing=opt.label_smoothing,
mirror=opt.mirror_loss,
fast_xentropy=opt.fast_xentropy)
# This function replaces modules with the more optimized counterparts so that it can run faster
# Currently exp with LayerNorm
if not opt.memory_profiling:
# distributed is required to convert BatchNorm to SyncBatchNorm for DDP
optimize_model(model, distributed=(self.world_size > 1))
init_model_parameters(model, opt)
self.model = model
self.loss_function = loss_function
# self.grad_scaler = torch.cuda.amp.GradScaler()
if self.cuda:
torch.cuda.set_device(self.device)
self.loss_function = self.loss_function.cuda(device=self.device)
self.model = self.model.cuda(device=self.device)
if opt.ctc_loss > 0.0:
self.ctc_loss_function = self.ctc_loss_function.cuda(
device=self.device)
if opt.load_from:
checkpoint = torch.load(opt.load_from,
map_location=lambda storage, loc: storage)
if setup_optimizer:
self.optim = onmt.Optim(opt)
self.optim.set_parameters(self.model.parameters())
if self.is_main():
print("[INFO] Optimizer: ", self.optim.optimizer)
if opt.load_from:
if 'optim' in checkpoint and checkpoint[
'optim'] is not None and not opt.reset_optim:
self.optim.load_state_dict(checkpoint['optim'])
if not self.opt.fp16:
opt_level = "O0"
keep_batchnorm_fp32 = False
elif self.opt.fp16_mixed:
opt_level = "O1"
keep_batchnorm_fp32 = None
else:
opt_level = "O2"
keep_batchnorm_fp32 = False
self.opt_level = opt_level
if self.cuda:
self.model, self.optim.optimizer = amp.initialize(
self.model,
self.optim.optimizer,
opt_level=opt_level,
keep_batchnorm_fp32=keep_batchnorm_fp32,
loss_scale="dynamic",
verbosity=1 if self.opt.verbose else 1)
if opt.load_from:
self.model.load_state_dict(checkpoint['model'])
if prec_opt is not None and hasattr(prec_opt, "fp16_mixed"):
# Only load amp information if the mode is the same
# Maybe its better to change between optimization mode?
if opt.fp16_mixed == prec_opt.fp16_mixed and opt.fp16 == prec_opt.fp16:
if 'amp' in checkpoint:
try:
amp.load_state_dict(checkpoint['amp'])
except Exception:
# loading the amp state can fail
pass
if self.world_size > 1:
# find_unused_parameters may be required for dropped layer (parameters that are not connected to
# any particular graph)
# find_unused_parameters = True
self.model = DDP(self.model,
delay_allreduce=True,
gradient_average=False)
print("[INFO] Process %d ready." % self.rank, flush=True)