def test_load_no_checkpoint(self):
with contextlib.redirect_stdout(StringIO()):
trainer, epoch_itr = get_trainer_and_epoch_itr(0, 150, 0, 0)
trainer.get_train_iterator = MagicMock(return_value=epoch_itr)
self.patches["os.path.isfile"].return_value = False
_, epoch_itr = checkpoint_utils.load_checkpoint(self.args_mock, trainer)
itr = epoch_itr.next_epoch_itr(shuffle=False)
self.assertEqual(epoch_itr.epoch, 1)
self.assertEqual(epoch_itr.iterations_in_epoch, 0)
self.assertEqual(next(itr)["net_input"]["src_tokens"][0].item(), 0)
def test_load_full_checkpoint(self):
with contextlib.redirect_stdout(StringIO()):
trainer, epoch_itr = get_trainer_and_epoch_itr(2, 150, 300, 150)
trainer.get_train_iterator = MagicMock(return_value=epoch_itr)
_, epoch_itr = checkpoint_utils.load_checkpoint(
self.args_mock, trainer)
itr = epoch_itr.next_epoch_itr(shuffle=False)
self.assertEqual(epoch_itr.epoch, 3)
self.assertEqual(epoch_itr.iterations_in_epoch, 0)
self.assertEqual(next(itr)['net_input']['src_tokens'][0].item(), 0)
def get_ready(args, ac='a'):
train_args = deepcopy(args)
if ac == 'a':
train_args.restore_file = args.actor_restore_file
train_args.task = args.actor_task
train_args.criterion = args.actor_criterion
train_args.save_interval_updates = args.actor_save_update
elif ac == 'c':
train_args.restore_file = args.critic_restore_file
train_args.task = args.critic_task
train_args.criterion = args.critic_criterion
train_args.save_interval_updates = args.critic_save_update
task = tasks.setup_task(train_args)
model = task.build_model(train_args)
criterion = task.build_criterion(train_args)
logger.info(model)
logger.info('model {}, criterion {}'.format(train_args.arch,
criterion.__class__.__name__))
logger.info('num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
if train_args.model_parallel_size == 1:
trainer = Trainer(train_args, task, model, criterion)
else:
trainer = MegatronTrainer(train_args, task, model, criterion)
logger.info('training on {} GPUs'.format(
train_args.distributed_world_size))
logger.info(
'max tokens per GPU = {} and max sentences per GPU = {}'.format(
train_args.max_tokens,
train_args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(
train_args, trainer)
return train_args, task, model, criterion, trainer, epoch_itr, extra_state
def test_finetune_from_model_args_conflict(self):
with contextlib.redirect_stdout(StringIO()):
trainer, epoch_itr = get_trainer_and_epoch_itr(1, 150, 0, 0)
trainer.get_train_iterator = MagicMock(return_value=epoch_itr)
for arg in [
'reset_optimizer', 'reset_lr_scheduler', 'reset_meters',
'reset_dataloader'
]:
with self.subTest(arg=arg):
args_mock = get_mock_args("/temp/checkpoint_pretrained.pt")
setattr(args_mock, arg, True)
with self.assertRaises(Exception) as context:
_, _ = checkpoint_utils.load_checkpoint(
args_mock, trainer)
self.assertTrue(
"--finetune-from-model can not be set together with either --reset-optimizer"
" or reset_lr_scheduler or reset_meters or reset_dataloader"
in str(context.exception))
def prepare_task(args, xla_device):
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=True, epoch=0)
# Build models and criteria to print some metadata
torch.manual_seed(args.seed)
model, criterion = task.build_model(args), task.build_criterion(args)
xm.master_print(model)
xm.master_print('| model {}, criterion {}'.format(
args.arch, criterion.__class__.__name__))
xm.master_print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad)))
model = model.to(xla_device)
trainer = Trainer(args, task, model, criterion, xla_device=xla_device)
lr = trainer.get_lr()
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
# we overwrite distributed args here to shard data using torch_xla's
# distributed training.
trainer.args.distributed_rank = xm.get_ordinal()
trainer.args.distributed_world_size = xm.xrt_world_size()
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
trainer.args.distributed_rank = 0
trainer.args.distributed_world_size = 1
trainer.meters_to_device(xla_device)
valid_subsets = args.valid_subset.split(',')
ordinal = xm.get_ordinal(defval=-1)
device_str = (
str(xla_device) if ordinal < 0 else
'{}/{}'.format(xla_device, ordinal)
)
return task, trainer, model, epoch_itr, lr, valid_subsets, device_str
def test_finetune_from_model(self):
with contextlib.redirect_stdout(StringIO()):
trainer, epoch_itr = get_trainer_and_epoch_itr(1, 150, 0, 0)
trainer.get_train_iterator = MagicMock(return_value=epoch_itr)
from_model_path = "/temp/checkpoint_pretrained.pt"
args_mock = get_mock_args(from_model_path)
args_mock.restore_file = "checkpoint_last.pt"
def mock_finetune_exist(path):
if path == from_model_path:
return True
else:
return False
self.patches[
'fairseq.file_io.PathManager.exists'].side_effect = mock_finetune_exist
_, _ = checkpoint_utils.load_checkpoint(args_mock, trainer)
checkpoint_path, reset_optimizer, reset_lr_scheduler, \
optimizer_overrides = trainer.load_checkpoint.call_args[0]
reset_meters = trainer.load_checkpoint.call_args[1]['reset_meters']
self.assertTrue(reset_optimizer)
self.assertTrue(reset_lr_scheduler)
self.assertTrue(reset_meters)
def prepare_task(args, devices):
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=True, epoch=0)
# Build models and criteria to print some metadata
model_parallel = dp.DataParallel(
lambda: task.build_model(args), device_ids=devices)
model, criterion = task.build_model(args), task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
del model, criterion
# Build trainers
trainers = {
device: Trainer(args, task, model, task.build_criterion(args), xla=True)
for device, model in zip(model_parallel.devices, model_parallel.models)
}
trainer = trainers[devices[0]]
lr = trainer.get_lr()
# TODO(taylanbil): for now, this next line is only creating the iterator.
# validate its behavior with the case where a checkpoint actually exists.
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
valid_subsets = args.valid_subset.split(',')
return task, trainers, model_parallel, epoch_itr, lr, valid_subsets
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
print(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr, filtered_maxpos_indices = checkpoint_utils.load_checkpoint(
args, trainer)
# pretrain data actor
# only the language actor model can be pretrained
if args.pretrain_laser and args.pretrain_data_actor and args.data_actor == 'ave':
# pretrain the agent with LASER score
# epoch_itr, indices = trainer.get_train_iterator(1)
path = '/home/wtan12/multiDDS/'
trainer.pretrain_LASER('en-ps.laser-score', epoch_itr)
if args.compare_laser:
epoch_itr, indices = trainer.get_train_iterator(1)
print('Number of Indices: ', len(indices))
scores = collections.defaultdict(float)
# compare with laser label using R^2 Score, only used after model is trained
# itr = epoch_itr.next_epoch_itr(fix_batches_to_gpus=False, shuffle=False)
data_actor = trainer.data_actor
itr = epoch_itr.next_epoch_itr(
fix_batches_to_gpus=args.fix_batches_to_gpus,
shuffle=False,
offset=0,
datasize=-1,
)
for i, sample in enumerate(itr):
sample = trainer._prepare_sample(sample)
sample = list(sample.values())[0]
score = data_actor(sample).cpu().detach().numpy().tolist()
indices = sample['id'].data.cpu().numpy().ravel().tolist()
for k, v in zip(indices, score):
scores[k] = float(v[0])
scores = sorted(scores.items(), key=lambda x: x[0])
print('Number of Indices in Scoring file: ', len(scores))
path = '/home/wtan12/multiDDS/'
with open(path + 'en-ps.laser-score', 'r') as r:
data = r.read()
laser_score = []
for i, item in enumerate(data.split('\n')):
laser_score.append(item)
laser_score.pop()
r2 = 0.0
with open(path + 'en-ps.dds_score', 'w') as f:
for k, v in scores:
f.write(str(v) + '\n')
truth = float(laser_score[k])
r2 += (truth - v)**2
print('R2 Score compared to LASER file: ', r2)
return
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
if args.eval_bleu:
generator = task.build_generator(args)
args.maximize_best_checkpoint_metric = True
else:
generator = None
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
epoch_itr = train(args, trainer, task, epoch_itr, generator,
filtered_maxpos_indices)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets, generator)
else:
valid_losses = [None]
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr,
valid_losses[0])
if ':' in getattr(args, 'data', ''):
# sharded data: get train iterator for next epoch
epoch_itr = trainer.get_train_iterator(epoch_itr.epoch)[0]
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args, config=None, init_distributed=False):
utils.import_user_module(args)
experiment = None
if config:
experiment = ExistingExperiment(
api_key=config["api_key"],
previous_experiment=config["experiment_key"],
auto_output_logging=None,
)
assert (
args.max_tokens is not None or args.max_sentences is not None
), "Must specify batch size either with --max-tokens or --max-sentences"
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
print(args)
if experiment:
experiment.log_parameters(vars(args),
prefix="Device {} :: ".format(
args.device_id))
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(","):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print("| model {}, criterion {}".format(args.arch,
criterion.__class__.__name__))
print("| num. model params: {} (num. trained: {})".format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
if experiment:
experiment.log_parameters(
{
"criterion":
criterion.__class__.__name__,
"num. model params":
sum(p.numel() for p in model.parameters()),
"num. trained params":
sum(p.numel() for p in model.parameters() if p.requires_grad),
},
prefix="Device {} :: ".format(args.device_id),
)
# Build trainer
trainer = Trainer(args, task, model, criterion)
print("| training on {} GPUs".format(args.distributed_world_size))
print("| max tokens per GPU = {} and max sentences per GPU = {}".format(
args.max_tokens, args.max_sentences))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(",")
while (lr > args.min_lr and epoch_itr.epoch < max_epoch
and trainer.get_num_updates() < max_update):
# train for one epoch
train(args, trainer, task, epoch_itr, experiment)
if (not args.disable_validation
and epoch_itr.epoch % args.validate_interval == 0):
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets, experiment)
else:
valid_losses = [None]
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr,
valid_losses[0])
reload_dataset = ":" in getattr(args, "data", "")
# sharded data: get train iterator for next epoch
epoch_itr = trainer.get_train_iterator(epoch_itr.epoch,
load_dataset=reload_dataset)
train_meter.stop()
print("| done training in {:.1f} seconds".format(train_meter.sum))
if experiment:
experiment.log_metrics(
{
"valid_loss": valid_losses[0],
"lr": lr
},
prefix="Device {} ".format(args.device_id),
)
def main(args, ITE=0):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
reinit = True if args.prune_type=="reinit" else False
# Data Loader
transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,))])
num_workers=0
if args.dataset == "mnist":
traindataset = datasets.MNIST('../data', train=True, download=True,transform=transform)
testdataset = datasets.MNIST('../data', train=False, transform=transform)
from archs.mnist import AlexNet, LeNet5, fc1, vgg, resnet
elif args.dataset == "cifar10":
traindataset = datasets.CIFAR10('../data', train=True, download=True,transform=transform)
testdataset = datasets.CIFAR10('../data', train=False, transform=transform)
from archs.cifar10 import AlexNet, LeNet5, fc1, vgg, mobilenet, resnet, densenet
elif args.dataset == "fashionmnist":
traindataset = datasets.FashionMNIST('../data', train=True, download=True,transform=transform)
testdataset = datasets.FashionMNIST('../data', train=False, transform=transform)
from archs.mnist import AlexNet, LeNet5, fc1, vgg, resnet
elif args.dataset == "cifar100":
traindataset = datasets.CIFAR100('../data', train=True, download=True,transform=transform)
testdataset = datasets.CIFAR100('../data', train=False, transform=transform)
from archs.cifar100 import AlexNet, fc1, LeNet5, vgg, resnet, mobilenet
elif args.dataset == "imagenet":
traindataset = datasets.ImageNet('../data', train=True, download=True,transform=transform)
testdataset = datasets.ImageNet('../data', train=False, transform=transform)
from archs.imagenet import AlexNet, fc1, LeNet5, vgg, resnet
elif args.dataset == "ncf":
train_data, test_data, user_num ,item_num, train_mat = ncf.data_utils.load_all()
traindataset = ncf.data_utils.NCFData(train_data, item_num, train_mat, args.num_ng, True)
testdataset = ncf.data_utils.NCFData(test_data, item_num, train_mat, 0, False)
num_workers=4
elif args.dataset == "iwslt14":
fairsequtils.import_user_module(args)
task = tasks.setup_task(args)
valid_subset='valid'
task.load_dataset(valid_subset, combine=False, epoch=0)
#model = task.build_model(args)
#criterion = task.build_criterion(args)
#trainer = Trainer(args, task, model, criterion)
traindataset = None
testdataset = None
else:
print("\nWrong Dataset choice \n")
exit()
if traindataset:
train_loader = torch.utils.data.DataLoader(traindataset, batch_size=args.batch_size, shuffle=True, num_workers=num_workers,drop_last=False)
#train_loader = cycle(train_loader)
test_loader = torch.utils.data.DataLoader(testdataset, batch_size=args.batch_size, shuffle=False, num_workers=0,drop_last=True)
# Importing Network Architecture
global model
if args.arch_type == "fc1":
model = fc1.fc1().to(device)
elif args.arch_type == "lenet5":
model = LeNet5.LeNet5().to(device)
elif args.arch_type == "alexnet":
model = AlexNet.AlexNet().to(device)
elif args.arch_type == "vgg16":
model = vgg.vgg16().to(device)
elif args.arch_type == "resnet50":
model = resnet.ResNet50().to(device)
elif args.arch_type == "mobilenet":
model = mobilenet.mobilenet_v2().to(device)
elif args.arch_type == "densenet121":
model = densenet.densenet121().to(device)
elif args.arch_type == "transformer":
model = task.build_model(args).to(device)
# If you want to add extra model paste here
elif args.arch_type == "neumf":
model = ncf.model.NCF(user_num, item_num, args.factor_num, args.num_layers,
args.dropout, ncf.config.model, None, None)
# Weight Initialization
model.apply(weight_init)
# Copying and Saving Initial State
initial_state_dict = copy.deepcopy(model.state_dict())
utils.checkdir(f"{os.getcwd()}/saves/{args.arch_type}/{args.dataset}/")
torch.save(model, f"{os.getcwd()}/saves/{args.arch_type}/{args.dataset}/initial_state_dict_{args.prune_type}.pth.tar")
# Making Initial Mask
make_mask(model)
# Optimizer and Loss
if args.arch_type == "neumf":
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.BCEWithLogitsLoss()
elif args.arch_type != "transformer":
optimizer = torch.optim.Adam(model.parameters(), weight_decay=1e-4)
criterion = nn.CrossEntropyLoss() # Default was F.nll_loss
else:
criterion = task.build_criterion(args)
trainer = ftrainer.Trainer(args, task, model, criterion)
assert(trainer.model)
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Layer Looper
for name, param in model.named_parameters():
print(name, param.size())
# Pruning
# NOTE First Pruning Iteration is of No Compression
bestacc = 0.0
best_accuracy = 0
ITERATION = args.prune_iterations
comp = np.zeros(ITERATION,float)
bestacc = np.zeros(ITERATION,float)
step = 0
all_loss = np.zeros(args.end_iter,float)
all_accuracy = np.zeros(args.end_iter,float)
for _ite in range(args.start_iter, ITERATION):
if not _ite == 0:
prune_by_percentile(args.prune_percent, resample=resample, reinit=reinit)
if reinit:
model.apply(weight_init)
#if args.arch_type == "fc1":
# model = fc1.fc1().to(device)
#elif args.arch_type == "lenet5":
# model = LeNet5.LeNet5().to(device)
#elif args.arch_type == "alexnet":
# model = AlexNet.AlexNet().to(device)
#elif args.arch_type == "vgg16":
# model = vgg.vgg16().to(device)
#elif args.arch_type == "resnet18":
# model = resnet.resnet18().to(device)
#elif args.arch_type == "densenet121":
# model = densenet.densenet121().to(device)
#else:
# print("\nWrong Model choice\n")
# exit()
step = 0
for name, param in model.named_parameters():
if 'weight' in name:
weight_dev = param.device
param.data = torch.from_numpy(param.data.cpu().numpy() * mask[step]).to(weight_dev)
step = step + 1
step = 0
else:
original_initialization(mask, initial_state_dict)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lrate, weight_decay=1e-4)
print(f"\n--- Pruning Level [{ITE}:{_ite}/{ITERATION}]: ---")
# Print the table of Nonzeros in each layer
comp1 = utils.print_nonzeros(model)
comp[_ite] = comp1
pbar = tqdm(range(args.end_iter))
for iter_ in pbar:
# Frequency for Testing
if iter_ % args.valid_freq == 0:
if args.arch_type == "neumf":
HR, NDCG = ncf.evaluate.metrics(model, test_loader, args.top_k)
accuracy = HR
elif args.arch_type != "transformer":
accuracy = test(model, test_loader, criterion)
else:
# Hack
#accuracy = best_accuracy + 1
accuracy = test_nmt(args, trainer, task, epoch_itr)
# Save Weights
if accuracy > best_accuracy:
best_accuracy = accuracy
utils.checkdir(f"{os.getcwd()}/saves/{args.arch_type}/{args.dataset}/")
torch.save(model,f"{os.getcwd()}/saves/{args.arch_type}/{args.dataset}/{_ite}_model_{args.prune_type}.pth.tar")
# Training
if args.arch_type == "neumf":
loss = train_neumf(model, train_loader, test_loader, optimizer, criterion)
elif args.arch_type != "transformer":
loss = train(model, train_loader, optimizer, criterion)
else:
loss = train_nmt(args, trainer, task, epoch_itr)
all_loss[iter_] = loss
all_accuracy[iter_] = accuracy
# Frequency for Printing Accuracy and Loss
if iter_ % args.print_freq == 0:
pbar.set_description(
f'Train Epoch: {iter_}/{args.end_iter} Loss: {loss:.6f} Accuracy: {accuracy:.2f}% Best Accuracy: {best_accuracy:.2f}%')
writer.add_scalar('Accuracy/test', best_accuracy, comp1)
bestacc[_ite]=best_accuracy
# Plotting Loss (Training), Accuracy (Testing), Iteration Curve
#NOTE Loss is computed for every iteration while Accuracy is computed only for every {args.valid_freq} iterations. Therefore Accuracy saved is constant during the uncomputed iterations.
#NOTE Normalized the accuracy to [0,100] for ease of plotting.
plt.plot(np.arange(1,(args.end_iter)+1), 100*(all_loss - np.min(all_loss))/np.ptp(all_loss).astype(float), c="blue", label="Loss")
plt.plot(np.arange(1,(args.end_iter)+1), all_accuracy, c="red", label="Accuracy")
plt.title(f"Loss Vs Accuracy Vs Iterations ({args.dataset},{args.arch_type})")
plt.xlabel("Iterations")
plt.ylabel("Loss and Accuracy")
plt.legend()
plt.grid(color="gray")
utils.checkdir(f"{os.getcwd()}/plots/lt/{args.arch_type}/{args.dataset}/")
plt.savefig(f"{os.getcwd()}/plots/lt/{args.arch_type}/{args.dataset}/{args.prune_type}_LossVsAccuracy_{comp1}.png", dpi=1200)
plt.close()
# Dump Plot values
utils.checkdir(f"{os.getcwd()}/dumps/lt/{args.arch_type}/{args.dataset}/")
all_loss.dump(f"{os.getcwd()}/dumps/lt/{args.arch_type}/{args.dataset}/{args.prune_type}_all_loss_{comp1}.dat")
all_accuracy.dump(f"{os.getcwd()}/dumps/lt/{args.arch_type}/{args.dataset}/{args.prune_type}_all_accuracy_{comp1}.dat")
# Dumping mask
utils.checkdir(f"{os.getcwd()}/dumps/lt/{args.arch_type}/{args.dataset}/")
with open(f"{os.getcwd()}/dumps/lt/{args.arch_type}/{args.dataset}/{args.prune_type}_mask_{comp1}.pkl", 'wb') as fp:
pickle.dump(mask, fp)
# Making variables into 0
best_accuracy = 0
all_loss = np.zeros(args.end_iter,float)
all_accuracy = np.zeros(args.end_iter,float)
# Dumping Values for Plotting
utils.checkdir(f"{os.getcwd()}/dumps/lt/{args.arch_type}/{args.dataset}/")
comp.dump(f"{os.getcwd()}/dumps/lt/{args.arch_type}/{args.dataset}/{args.prune_type}_compression.dat")
bestacc.dump(f"{os.getcwd()}/dumps/lt/{args.arch_type}/{args.dataset}/{args.prune_type}_bestaccuracy.dat")
# Plotting
a = np.arange(args.prune_iterations)
plt.plot(a, bestacc, c="blue", label="Winning tickets")
plt.title(f"Test Accuracy vs Unpruned Weights Percentage ({args.dataset},{args.arch_type})")
plt.xlabel("Unpruned Weights Percentage")
plt.ylabel("test accuracy")
plt.xticks(a, comp, rotation ="vertical")
plt.ylim(0,100)
plt.legend()
plt.grid(color="gray")
utils.checkdir(f"{os.getcwd()}/plots/lt/{args.arch_type}/{args.dataset}/")
plt.savefig(f"{os.getcwd()}/plots/lt/{args.arch_type}/{args.dataset}/{args.prune_type}_AccuracyVsWeights.png", dpi=1200)
plt.close()
def main(cfg: FairseqConfig) -> None:
if isinstance(cfg, argparse.Namespace):
cfg = convert_namespace_to_omegaconf(cfg)
utils.import_user_module(cfg.common)
if is_master(cfg.distributed_training) and "job_logging_cfg" in cfg:
# make hydra logging work with ddp (see # see https://github.com/facebookresearch/hydra/issues/1126)
logging.config.dictConfig(OmegaConf.to_container(cfg.job_logging_cfg))
assert (
cfg.dataset.max_tokens is not None
or cfg.dataset.batch_size is not None
), "Must specify batch size either with --max-tokens or --batch-size"
metrics.reset()
np.random.seed(cfg.common.seed)
utils.set_torch_seed(cfg.common.seed)
if distributed_utils.is_master(cfg.distributed_training):
checkpoint_utils.verify_checkpoint_directory(cfg.checkpoint.save_dir)
# Print args
logger.info(cfg)
if cfg.checkpoint.write_checkpoints_asynchronously:
try:
import iopath # noqa: F401
except ImportError:
logging.exception(
"Asynchronous checkpoint writing is specified but iopath is "
"not installed: `pip install iopath`")
return
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(cfg.task)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in cfg.dataset.valid_subset.split(","):
task.load_dataset(valid_sub_split, combine=False, epoch=1)
assert cfg.criterion, "Please specify criterion to train a model"
# Build model and criterion
model = task.build_model(cfg.model)
criterion = task.build_criterion(cfg.criterion)
logger.info(model)
logger.info("task: {}".format(task.__class__.__name__))
logger.info("model: {}".format(model.__class__.__name__))
logger.info("criterion: {}".format(criterion.__class__.__name__))
logger.info("num. model params: {:,} (num. trained: {:,})".format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# (optionally) Configure quantization
if cfg.common.quantization_config_path is not None:
quantizer = quantization_utils.Quantizer(
config_path=cfg.common.quantization_config_path,
max_epoch=cfg.optimization.max_epoch,
max_update=cfg.optimization.max_update,
)
else:
quantizer = None
# Build trainer
if cfg.common.model_parallel_size == 1:
trainer = Trainer(cfg, task, model, criterion, quantizer)
else:
trainer = MegatronTrainer(cfg, task, model, criterion)
logger.info("training on {} devices (GPUs/TPUs)".format(
cfg.distributed_training.distributed_world_size))
logger.info("max tokens per GPU = {} and batch size per GPU = {}".format(
cfg.dataset.max_tokens,
cfg.dataset.batch_size,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(
cfg.checkpoint,
trainer,
# don't cache epoch iterators for sharded datasets
disable_iterator_cache=task.has_sharded_data("train"),
)
max_epoch = cfg.optimization.max_epoch or math.inf
lr = trainer.get_lr()
train_meter = meters.StopwatchMeter()
train_meter.start()
while epoch_itr.next_epoch_idx <= max_epoch:
if lr <= cfg.optimization.stop_min_lr:
logger.info(
f"stopping training because current learning rate ({lr}) is smaller "
"than or equal to minimum learning rate "
f"(--stop-min-lr={cfg.optimization.stop_min_lr})")
break
# train for one epoch
valid_losses, should_stop = train(cfg, trainer, task, epoch_itr)
if should_stop:
break
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
epoch_itr = trainer.get_train_iterator(
epoch_itr.next_epoch_idx,
# sharded data: get train iterator for next epoch
load_dataset=task.has_sharded_data("train"),
# don't cache epoch iterators for sharded datasets
disable_iterator_cache=task.has_sharded_data("train"),
)
train_meter.stop()
logger.info("done training in {:.1f} seconds".format(train_meter.sum))
# ioPath implementation to wait for all asynchronous file writes to complete.
if cfg.checkpoint.write_checkpoints_asynchronously:
logger.info(
"ioPath PathManager waiting for all asynchronous checkpoint "
"writes to finish.")
PathManager.async_close()
logger.info("ioPath PathManager finished waiting.")
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
logger.info(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
logger.info(model)
logger.info('model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
logger.info('num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = Trainer(args, task, model, criterion)
logger.info('training on {} GPUs'.format(args.distributed_world_size))
logger.info('max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
print(args.multi_views)
while (
lr > args.min_lr
and (
epoch_itr.epoch < max_epoch
# allow resuming training from the final checkpoint
or epoch_itr._next_epoch_itr is not None
)
and trainer.get_num_updates() < max_update
):
# train for one epoch
train(args, trainer, task, epoch_itr)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr, valid_subsets)
else:
valid_losses = [None]
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
bart = BARTHubInterface(args, task, trainer.model).cuda()
#print(bart.device)
bart.eval()
count = 1
bsz = 8
print("Test on val set: ")
with open('../data/val_sent_trans_cons_label.source') as source, open('../data/val_sent_c99_label.source') as source2, open('./val_best_multi_attn_'+str(args.lr_weight)+'_.hypo', 'wt', encoding='utf-8') as fout:
s1 = source.readlines()
s2 = source2.readlines()
slines = [s1[0].strip()]
slines2 = [s2[0].strip()]
for i in tqdm(range(1, len(s1))):
if count % bsz == 0:
with torch.no_grad():
if args.multi_views:
hypotheses_batch = bart.sample(slines, sentences2 = slines2, balance = True, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
else:
hypotheses_batch = bart.sample(slines, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
for hypothesis in hypotheses_batch:
fout.write(hypothesis + '\n')
fout.flush()
slines = []
slines2 = []
slines.append(s1[i].strip())
slines2.append(s2[i].strip())
count += 1
if slines != []:
if args.multi_views:
hypotheses_batch = bart.sample(slines, sentences2 = slines2, balance = True, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
else:
hypotheses_batch = bart.sample(slines, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
#hypotheses_batch = bart.sample(slines, sentences2 = slines2, balance = True, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
for hypothesis in hypotheses_batch:
fout.write(hypothesis + '\n')
fout.flush()
hyp_path = './val_best_multi_attn_'+str(args.lr_weight)+'_.hypo'
ref_path = '../data/val_sent_trans_cons_label.target'
hypothesis = []
with open(hyp_path, 'r') as f:
lines = f.readlines()
for l in lines:
hypothesis.append(l[:-1])
reference = []
with open(ref_path, 'r') as f:
lines = f.readlines()
for l in lines:
reference.append(l[:-1])
rouge = Rouge()
print("Val", rouge.get_scores(hypothesis, reference, avg = True))
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
print("Test on testing set: ")
count = 1
bsz = 8
with open('../data/test_sent_trans_cons_label.source') as source, open('../data/test_sent_c99_label.source') as source2, open('./test_best_multi_attn_'+str(args.lr_weight)+'_.hypo', 'wt', encoding='utf-8') as fout:
s1 = source.readlines()
s2 = source2.readlines()
slines = [s1[0].strip()]
slines2 = [s2[0].strip()]
for i in tqdm(range(1, len(s1))):
if count % bsz == 0:
with torch.no_grad():
if args.multi_views:
hypotheses_batch = bart.sample(slines, sentences2 = slines2, balance = True, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
else:
hypotheses_batch = bart.sample(slines, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
for hypothesis in hypotheses_batch:
fout.write(hypothesis + '\n')
fout.flush()
slines = []
slines2 = []
slines.append(s1[i].strip())
slines2.append(s2[i].strip())
count += 1
if slines != []:
if args.multi_views:
hypotheses_batch = bart.sample(slines, sentences2 = slines2, balance = True, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
else:
hypotheses_batch = bart.sample(slines, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
for hypothesis in hypotheses_batch:
fout.write(hypothesis + '\n')
fout.flush()
hyp_path = './test_best_multi_attn_'+str(args.lr_weight)+'_.hypo'
ref_path = '../data/test_sent_trans_cons_label.target'
hypothesis = []
with open(hyp_path, 'r') as f:
lines = f.readlines()
for l in lines:
hypothesis.append(l[:-1])
reference = []
with open(ref_path, 'r') as f:
lines = f.readlines()
for l in lines:
reference.append(l[:-1])
rouge = Rouge()
print('Test', rouge.get_scores(hypothesis, reference, avg = True))
# early stop
if should_stop_early(args, valid_losses[0]):
logger.info('early stop since valid performance hasn\'t improved for last {} runs'.format(args.patience))
break
epoch_itr = trainer.get_train_iterator(
epoch_itr.epoch,
# sharded data: get train iterator for next epoch
load_dataset=(os.pathsep in getattr(args, 'data', '')),
)
train_meter.stop()
logger.info('done training in {:.1f} seconds'.format(train_meter.sum))
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
logger.info(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
logger.info(model)
logger.info('model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
logger.info('num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = Trainer(args, task, model, criterion)
logger.info('training on {} GPUs'.format(args.distributed_world_size))
logger.info(
'max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
tokenize = sacrebleu.DEFAULT_TOKENIZER if not args.eval_tokenized_bleu else 'none'
hyps, refs = validate(args, trainer, task, epoch_itr, valid_subsets)
for h, r, split in zip(hyps, refs, args.valid_subset.split(',')):
assert len(h) == len(r)
sacrebleu_score, _, _ = sacrebleu.corpus_bleu(
h, [r], tokenize=tokenize), hyps, refs
bleu = compute_cvpr_bleu(h, r)
rouge_score = rouge.rouge(h, r)
print('{} set has {} samples,\n'
'sacrebleu: {},\n'
'CVPR BLEU scripts: {}\n'
'CVPR ROUGE: {}'.format(split, len(h), sacrebleu_score, bleu,
rouge_score))
print('performance: {:.2f} {}'.format(
rouge_score['rouge_l/f_score'] * 100,
' '.join([str(b) for b in bleu])))
def main(args):
import_user_module(args)
assert (
args.max_tokens is not None or args.batch_size is not None
), "Must specify batch size either with --max-tokens or --batch-size"
metrics.reset()
np.random.seed(args.seed)
utils.set_torch_seed(args.seed)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
logger.info(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(","):
task.load_dataset(valid_sub_split, combine=False, epoch=1)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
logger.info(model)
logger.info("task: {} ({})".format(args.task, task.__class__.__name__))
logger.info("model: {} ({})".format(args.arch, model.__class__.__name__))
logger.info("criterion: {} ({})".format(args.criterion,
criterion.__class__.__name__))
logger.info("num. model params: {} (num. trained: {})".format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# breakpoint()
# ========== initialize the model with pretrained BART parameters ==========
# for shared embeddings and subtoken split for amr nodes
if 'bartsv' in args.arch:
if args.initialize_with_bart:
logger.info(
'-' * 10 +
' initializing model parameters with pretrained BART model ' +
'-' * 10)
new_state_dict = copy.deepcopy(task.bart.model.state_dict())
# treat the embedding initialization separately later, as the size different
logger.info(
'-' * 10 +
' delay encoder embeddings, decoder input and output embeddings initialization '
+ '-' * 10)
ignore_keys = set([
'encoder.embed_tokens.weight', 'decoder.embed_tokens.weight',
'decoder.output_projection.weight'
])
for k in ignore_keys:
del new_state_dict[k]
if not args.initialize_with_bart_enc:
logger.info(
'-' * 10 +
' do not initialize with BART encoder parameters ' +
'-' * 10)
for k in list(new_state_dict.keys()):
if k.startswith('encoder'):
del new_state_dict[k]
if not args.initialize_with_bart_dec:
logger.info(
'-' * 10 +
' do not initialize with BART decoder parameters ' +
'-' * 10)
for k in list(new_state_dict.keys()):
if k.startswith('decoder'):
del new_state_dict[k]
model.load_state_dict(new_state_dict, strict=False, args=args)
# initialize the Bart part embeddings
bart_vocab_size = task.target_dictionary.bart_vocab_size
# NOTE we need to prune the pretrained BART embeddings, especially for bart.base
bart_embed_weight = task.bart.model.encoder.embed_tokens.weight.data[:
bart_vocab_size]
assert len(bart_embed_weight) == bart_vocab_size
with torch.no_grad():
model.encoder.embed_tokens.weight[:bart_vocab_size].copy_(
bart_embed_weight)
model.decoder.embed_tokens.weight[:bart_vocab_size].copy_(
bart_embed_weight)
model.decoder.output_projection.weight[:bart_vocab_size].copy_(
bart_embed_weight)
if args.bart_emb_init_composition:
logger.info(
'-' * 10 +
' initialize extended target embeddings with compositional embeddings '
'from BART vocabulary ' + '-' * 10)
# breakpoint()
symbols = [
task.target_dictionary[idx]
for idx in range(bart_vocab_size, len(task.target_dictionary))
]
mapper = MapAvgEmbeddingBART(task.bart,
task.bart.model.decoder.embed_tokens)
comp_embed_weight, map_all = mapper.map_avg_embeddings(
symbols, transform=transform_action_symbol, add_noise=False)
assert len(comp_embed_weight) == len(symbols)
with torch.no_grad():
model.encoder.embed_tokens.weight[bart_vocab_size:].copy_(
comp_embed_weight)
model.decoder.embed_tokens.weight[bart_vocab_size:].copy_(
comp_embed_weight)
model.decoder.output_projection.weight[bart_vocab_size:].copy_(
comp_embed_weight)
elif 'bart' in args.arch:
if args.initialize_with_bart:
logger.info(
'-' * 10 +
' initializing model parameters with pretrained BART model ' +
'-' * 10)
new_state_dict = copy.deepcopy(task.bart.model.state_dict())
if not args.bart_emb_decoder:
logger.info('-' * 10 +
' build a separate decoder dictionary embedding ' +
'-' * 10)
if not args.bart_emb_decoder_input:
ignore_keys = set([
'decoder.embed_tokens.weight',
'decoder.output_projection.weight'
])
else:
logger.info(
'-' * 10 +
' use BART dictionary embedding for target input ' +
'-' * 10)
ignore_keys = set(['decoder.output_projection.weight'])
for k in ignore_keys:
del new_state_dict[k]
if not args.initialize_with_bart_enc:
logger.info(
'-' * 10 +
' do not initialize with BART encoder parameters ' +
'-' * 10)
for k in list(new_state_dict.keys()):
if k.startswith('encoder'):
del new_state_dict[k]
if not args.initialize_with_bart_dec:
logger.info(
'-' * 10 +
' do not initialize with BART decoder parameters ' +
'-' * 10)
for k in list(new_state_dict.keys()):
if k.startswith('decoder'):
del new_state_dict[k]
model.load_state_dict(new_state_dict, strict=False, args=args)
# initialize the target embeddings with average of subtoken embeddings in BART vocabulary
if args.bart_emb_init_composition:
assert not args.bart_emb_decoder, 'should not use the compositional embeddings on top of BART vocabulary here'
logger.info(
'-' * 10 +
' initialize target embeddings with compositional embeddings from BART vocabulary '
+ '-' * 10)
composite_embed = CompositeEmbeddingBART(
task.bart, task.bart.model.decoder.embed_tokens,
task.target_dictionary)
if args.bart_emb_decoder_input:
# only initialize the decoder output embeddings
with torch.no_grad():
model.decoder.output_projection.weight.copy_(
composite_embed.embedding_weight)
else:
# initialize both the decoder input and output embeddings
with torch.no_grad():
model.decoder.embed_tokens.weight.copy_(
composite_embed.embedding_weight)
model.decoder.output_projection.weight.copy_(
composite_embed.embedding_weight)
elif 'roberta' in args.arch:
# initialize the target embeddings with average of subtoken embeddings in BART vocabulary
if args.bart_emb_init_composition:
assert not args.bart_emb_decoder, 'should not use the compositional embeddings on top of RoBERTa vocabulary here'
logger.info(
'-' * 10 +
' initialize target embeddings with compositional embeddings from RoBERTa vocabulary '
+ '-' * 10)
composite_embed = CompositeEmbeddingBART(
task.bart, # NOTE here "bart" means roberta
task.bart.model.encoder.sentence_encoder.embed_tokens,
task.target_dictionary)
if args.bart_emb_decoder_input:
# only initialize the decoder output embeddings
with torch.no_grad():
model.decoder.output_projection.weight.copy_(
composite_embed.embedding_weight)
else:
# initialize both the decoder input and output embeddings
with torch.no_grad():
model.decoder.embed_tokens.weight.copy_(
composite_embed.embedding_weight)
model.decoder.output_projection.weight.copy_(
composite_embed.embedding_weight)
else:
raise ValueError
# ==========================================================================
# breakpoint()
# (optionally) Configure quantization
if args.quantization_config_path is not None:
quantizer = quantization_utils.Quantizer(
config_path=args.quantization_config_path,
max_epoch=args.max_epoch,
max_update=args.max_update,
)
else:
quantizer = None
# Build trainer
if args.model_parallel_size == 1:
trainer = Trainer(args, task, model, criterion, quantizer)
else:
trainer = MegatronTrainer(args, task, model, criterion)
logger.info("training on {} devices (GPUs/TPUs)".format(
args.distributed_world_size))
logger.info(
"max tokens per GPU = {} and max sentences per GPU = {}".format(
args.max_tokens, args.batch_size))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(
args,
trainer,
# don't cache epoch iterators for sharded datasets
disable_iterator_cache=task.has_sharded_data("train"),
)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
lr = trainer.get_lr()
train_meter = meters.StopwatchMeter()
train_meter.start()
while lr > args.min_lr and epoch_itr.next_epoch_idx <= max_epoch:
# train for one epoch
valid_losses, should_stop = train(args, trainer, task, epoch_itr)
if should_stop:
break
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
epoch_itr = trainer.get_train_iterator(
epoch_itr.next_epoch_idx,
# sharded data: get train iterator for next epoch
load_dataset=task.has_sharded_data("train"),
# don't cache epoch iterators for sharded datasets
disable_iterator_cache=task.has_sharded_data("train"),
)
train_meter.stop()
logger.info("done training in {:.1f} seconds".format(train_meter.sum))
def main(args):
utils.import_user_module(args)
utils.handle_save_path(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Print args
print(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
# Build trainer
trainer = Trainer(args, task, model, criterion)
# Load the latest checkpoint if one is available and restore the corresponding train iterator
args.train_subset = 'valid' # no need to train, so just set a small subset to save loading time
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# run evolutionary search to find the model with lowest loss and satisfies the latency requirement
evolver = Evolution(args, trainer, task, epoch_itr)
best_config = evolver.run_evo_search()
with open(args.write_config_path, 'w') as fid:
encoder_layer_num = best_config['encoder']['encoder_layer_num']
decoder_layer_num = best_config['decoder']['decoder_layer_num']
fid.write(
f"encoder-embed-dim-subtransformer: {best_config['encoder']['encoder_embed_dim']}\n"
)
fid.write(
f"decoder-embed-dim-subtransformer: {best_config['decoder']['decoder_embed_dim']}\n\n"
)
fid.write(
f"encoder-ffn-embed-dim-all-subtransformer: {best_config['encoder']['encoder_ffn_embed_dim'][:encoder_layer_num]}\n"
)
fid.write(
f"decoder-ffn-embed-dim-all-subtransformer: {best_config['decoder']['decoder_ffn_embed_dim'][:decoder_layer_num]}\n\n"
)
fid.write(
f"encoder-layer-num-subtransformer: {best_config['encoder']['encoder_layer_num']}\n"
)
fid.write(
f"decoder-layer-num-subtransformer: {best_config['decoder']['decoder_layer_num']}\n\n"
)
fid.write(
f"encoder-self-attention-heads-all-subtransformer: {best_config['encoder']['encoder_self_attention_heads'][:encoder_layer_num]}\n"
)
fid.write(
f"decoder-self-attention-heads-all-subtransformer: {best_config['decoder']['decoder_self_attention_heads'][:decoder_layer_num]}\n"
)
fid.write(
f"decoder-ende-attention-heads-all-subtransformer: {best_config['decoder']['decoder_ende_attention_heads'][:decoder_layer_num]}\n\n"
)
fid.write(
f"decoder-arbitrary-ende-attn-all-subtransformer: {best_config['decoder']['decoder_arbitrary_ende_attn'][:decoder_layer_num]}\n\n"
)
def main(args, init_distributed=False):
utils.import_user_module(args)
try:
from fairseq.fb_pathmgr import fb_pathmgr
global fb_pathmgr_registerd
if not fb_pathmgr_registerd:
fb_pathmgr.register()
fb_pathmgr_registerd = True
except (ModuleNotFoundError, ImportError):
pass
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
print(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# filter the params that is unused for finetuing, ad-hoc for finetuing, should turn off when bert pretraining.
for n, p in model.named_parameters():
if "lm_head" in n:
p.requires_grad = False
# print(n)
# print(n, p.requires_grad, p.shape)
# for i, (n, p) in enumerate(model.named_parameters()):
# print(i, n, p.size())
# asdf
# Build trainer
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
if not hasattr(checkpoint_utils.save_checkpoint, 'not_best'):
checkpoint_utils.save_checkpoint.not_best = 0
#import pdb; pdb.set_trace()
while epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
print('Start training')
# train for one epoch
train(args, trainer, task, epoch_itr)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
if args.early_stop > 0:
if hasattr(
checkpoint_utils.save_checkpoint, 'best'
) and valid_losses[0] > checkpoint_utils.save_checkpoint.best:
checkpoint_utils.save_checkpoint.not_best += 1
print("| Not the best ckpt... not best:",
checkpoint_utils.save_checkpoint.not_best)
if checkpoint_utils.save_checkpoint.not_best > args.early_stop:
print("| Early stop...")
break
else:
checkpoint_utils.save_checkpoint.not_best = 0
else:
valid_losses = [None]
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr,
valid_losses[0])
reload_dataset = ':' in getattr(args, 'data', '')
# sharded data: get train iterator for next epoch
epoch_itr = trainer.get_train_iterator(epoch_itr.epoch,
load_dataset=reload_dataset)
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(cfg: DictConfig) -> None:
if isinstance(cfg, argparse.Namespace):
cfg = convert_namespace_to_omegaconf(cfg)
utils.import_user_module(cfg.common)
assert (
cfg.dataset.max_tokens is not None
or cfg.dataset.batch_size is not None
), "Must specify batch size either with --max-tokens or --batch-size"
metrics.reset()
np.random.seed(cfg.common.seed)
utils.set_torch_seed(cfg.common.seed)
if distributed_utils.is_master(cfg.distributed_training):
checkpoint_utils.verify_checkpoint_directory(cfg.checkpoint.save_dir)
# Print args
logger.info(cfg)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(cfg.task)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in cfg.dataset.valid_subset.split(","):
task.load_dataset(valid_sub_split, combine=False, epoch=1)
assert cfg.criterion, "Please specify criterion to train a model"
# Build model and criterion
model = task.build_model(cfg.model)
criterion = task.build_criterion(cfg.criterion)
logger.info(model)
logger.info("task: {}".format(task.__class__.__name__))
logger.info("model: {}".format(model.__class__.__name__))
logger.info("criterion: {}".format(criterion.__class__.__name__))
logger.info("num. model params: {} (num. trained: {})".format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# (optionally) Configure quantization
if cfg.common.quantization_config_path is not None:
quantizer = quantization_utils.Quantizer(
config_path=cfg.common.quantization_config_path,
max_epoch=cfg.optimization.max_epoch,
max_update=cfg.optimization.max_update,
)
else:
quantizer = None
# Build trainer
if cfg.common.model_parallel_size == 1:
trainer = Trainer(cfg, task, model, criterion, quantizer)
else:
trainer = MegatronTrainer(cfg, task, model, criterion)
logger.info("training on {} devices (GPUs/TPUs)".format(
cfg.distributed_training.distributed_world_size))
logger.info("max tokens per GPU = {} and batch size per GPU = {}".format(
cfg.dataset.max_tokens,
cfg.dataset.batch_size,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(
cfg.checkpoint,
trainer,
# don't cache epoch iterators for sharded datasets
disable_iterator_cache=task.has_sharded_data("train"),
)
# scores, sparsities, all_head_masks = mask_heads(cfg, task, trainer, epoch_itr, model, exact_pruning=True)
for k in [28, 32, 36]:
gibbs_sampling(cfg,
task,
trainer,
epoch_itr,
model,
early_stop_step=12,
K=k)
def main(args, init_distributed=False):
utils.import_user_module(args)
utils.handle_save_path(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
#if torch.cuda.is_available() and not args.cpu:
# torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
print(f"| Configs: {args}")
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(
f"| Model: {args.arch} \n| Criterion: {criterion.__class__.__name__}")
# Log architecture
if args.train_subtransformer:
print(" \n\n\t\tWARNING!!! Training one single SubTransformer\n\n")
print(
f"| SubTransformer Arch: {utils.get_subtransformer_config(args)} \n"
)
else:
print(" \n\n\t\tWARNING!!! Training SuperTransformer\n\n")
print(f"| SuperTransformer Arch: {model} \n")
# Log model size
if args.train_subtransformer:
print(
f"| SubTransformer size (without embedding weights): {model.get_sampled_params_numel(utils.get_subtransformer_config(args))}"
)
embed_size = args.decoder_embed_dim_subtransformer * len(task.tgt_dict)
print(f"| Embedding layer size: {embed_size} \n")
else:
model_s = 0
# if use model.state_dict, then will add 2 more parameters, they are encoder.version and decoder.version. Should not count them
for name, param in model.named_parameters():
if 'embed' not in name:
model_s += param.numel()
print(
f"| SuperTransofmer model size (without embedding weights): {model_s}"
)
print(
f"| Embedding layer size: {sum(p.numel() for p in model.parameters() if p.requires_grad) - model_s} \n"
)
# specify the length of the dummy input for profile
# for iwslt, the average length is 23, for wmt, that is 30
dummy_sentence_length_dict = {'iwslt': 23, 'wmt': 30}
if 'iwslt' in args.arch:
dummy_sentence_length = dummy_sentence_length_dict['iwslt']
elif 'wmt' in args.arch:
dummy_sentence_length = dummy_sentence_length_dict['wmt']
else:
raise NotImplementedError
dummy_src_tokens = [2] + [7] * (dummy_sentence_length - 1)
dummy_prev = [7] * (dummy_sentence_length - 1) + [2]
# profile the overall FLOPs number
if args.profile_flops:
import torchprofile
config_subtransformer = utils.get_subtransformer_config(args)
model.set_sample_config(config_subtransformer)
model.profile(mode=True)
macs = torchprofile.profile_macs(model,
args=(torch.tensor([dummy_src_tokens],
dtype=torch.long),
torch.tensor([30]),
torch.tensor([dummy_prev],
dtype=torch.long)))
model.profile(mode=False)
last_layer_macs = config_subtransformer['decoder'][
'decoder_embed_dim'] * dummy_sentence_length * len(task.tgt_dict)
print(f"| Total FLOPs: {macs * 2}")
print(f"| Last layer FLOPs: {last_layer_macs * 2}")
print(
f"| Total FLOPs without last layer: {(macs - last_layer_macs) * 2} \n"
)
exit(0)
with torch.autograd.set_detect_anomaly(True):
# Build trainer
trainer = Trainer(args, task, model, criterion)
print(f"| Training on {args.distributed_world_size} GPUs")
# print(f"| Max tokens per GPU = {args.max_tokens} and max sentences per GPU = {args.max_sentences} \n")
print(
f"| Max tokens per GPU = {args.max_tokens} and max sentences per GPU = {None} \n"
)
# Measure model latency, the program will exit after profiling latency
if args.latcpu or args.latgpu:
utils.measure_latency(args, model, dummy_src_tokens, dummy_prev)
exit(0)
# Load the latest checkpoint if one is available and restore the corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Evaluate the SubTransformer
if args.validate_subtransformer:
config = utils.get_subtransformer_config(args)
trainer.set_sample_config(config)
valid_loss = validate(args, trainer, task, epoch_itr, ['valid'],
'SubTransformer')
print(f"| SubTransformer validation loss:{valid_loss}")
# Loop boundaries
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
represent_configs = utils.get_represent_configs(args)
# Main training loop
while lr > args.stop_min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
for k, v in represent_configs.items():
trainer.set_sample_config(config=v)
valid_losses = validate(args,
trainer,
task,
epoch_itr,
valid_subsets,
sampled_arch_name=k)
else:
valid_losses = [None]
# update the best loss and get current lr; the real lr scheduling is done in trainer.train_step()
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint epoch level
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr,
valid_losses[0])
train_meter.stop()
print('| Done training in {:.1f} seconds'.format(train_meter.sum))
inds, coords, inds_nums = scheduler.get_job_indices(
losses, xcoordinates, ycoordinates, None)
start_time = time.time()
total_sync = 0.0
xtra_state = trainer.load_checkpoint(
os.path.join(args_transformer.save_dir, args_transformer.restore_file),
args_transformer.reset_optimizer,
args_transformer.reset_lr_scheduler,
eval(args_transformer.optimizer_overrides),
reset_meters=args_transformer.reset_meters,
)
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(
args_transformer, trainer)
update_freq = args_transformer.update_freq[epoch_itr.epoch - 1] if epoch_itr.epoch <= len(args_transformer.update_freq) else \
args_transformer.update_freq[-1]
# Initialize data iterator
itr = epoch_itr.next_epoch_itr(
fix_batches_to_gpus=args_transformer.fix_batches_to_gpus,
shuffle=(epoch_itr.epoch >= args_transformer.curriculum),
)
itr = iterators.GroupedIterator(itr, update_freq)
progress = progress_bar.build_progress_bar(
args_transformer,
itr,
epoch_itr.epoch,
no_progress_bar='simple',
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
logger.info(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=1)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
logger.info(model)
logger.info('model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
logger.info('num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = Trainer(args, task, model, criterion)
logger.info('training on {} GPUs'.format(args.distributed_world_size))
logger.info(
'max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = meters.StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
if args.eval_mode != 'none':
start_val_time = time.time()
with torch.no_grad():
if args.eval_mode != 'entropy':
_ = validate(args, trainer, task, epoch_itr, valid_subsets,
args.prune_num)
print('elapsed time (seconds): {}'.format(time.time() -
start_val_time))
_ = validate_iw(args,
trainer,
task,
epoch_itr,
valid_subsets,
args.prune_num,
mode=args.eval_mode)
return
while (lr > args.min_lr and epoch_itr.next_epoch_idx <= max_epoch
and trainer.get_num_updates() < max_update):
# train for one epoch
train(args, trainer, task, epoch_itr)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
else:
valid_losses = [None]
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr,
valid_losses[0])
# early stop
if should_stop_early(args, valid_losses[0]):
logger.info(
'early stop since valid performance hasn\'t improved for last {} runs'
.format(args.patience))
break
epoch_itr = trainer.get_train_iterator(
epoch_itr.next_epoch_idx,
# sharded data: get train iterator for next epoch
load_dataset=(os.pathsep in getattr(args, 'data', '')),
)
logger.info('done training in {:.1f} seconds'.format(train_meter.sum))
# _ = validate_iw(args, trainer, task, epoch_itr, valid_subsets)
train_meter.stop()
def main(args):
utils.import_user_module(args)
assert (
args.max_tokens is not None or args.max_sentences is not None
), "Must specify batch size either with --max-tokens or --max-sentences"
metrics.reset()
np.random.seed(args.seed)
utils.set_torch_seed(args.seed)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
checkpoint_utils.verify_checkpoint_directory(args.jason_log_dir)
# Print args
logger.info(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(","):
task.load_dataset(valid_sub_split, combine=False, epoch=1)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
logger.info(model)
logger.info("task: {} ({})".format(args.task, task.__class__.__name__))
logger.info("model: {} ({})".format(args.arch, model.__class__.__name__))
logger.info(
"criterion: {} ({})".format(args.criterion, criterion.__class__.__name__)
)
logger.info(
"num. model params: {} (num. trained: {})".format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
)
)
# (optionally) Configure quantization
if args.quantization_config_path is not None:
quantizer = quantization_utils.Quantizer(
config_path=args.quantization_config_path,
max_epoch=args.max_epoch,
max_update=args.max_update,
)
else:
quantizer = None
# Build trainer
if args.model_parallel_size == 1:
trainer = Trainer(args, task, model, criterion, quantizer)
else:
trainer = MegatronTrainer(args, task, model, criterion)
logger.info(
"training on {} devices (GPUs/TPUs)".format(args.distributed_world_size)
)
logger.info(
"max tokens per GPU = {} and max sentences per GPU = {}".format(
args.max_tokens, args.max_sentences
)
)
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(
args,
trainer,
# don't cache epoch iterators for sharded datasets
disable_iterator_cache=task.has_sharded_data("train"),
)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
lr = trainer.get_lr()
train_meter = meters.StopwatchMeter()
train_meter.start()
##### begin jason #####
updates_list = []; train_ppl_list = []; train_loss_list = []; val_ppl_list = []; val_loss_list = []; train_uid_loss_list = []; val_uid_loss_list = []
log_writer = open(os.path.join(args.save_dir, 'train_logs.csv'), 'w')
log_writer.write(f'updates,train_loss,train_ppl,val_loss,val_ppl\n')
backup_writefile = os.path.join(args.jason_log_dir, 'train_logs_backup.csv')
os.system(f'touch {backup_writefile}')
os.system(f'echo "updates,train_loss,train_ppl,val_loss,val_ppl,train_uid_loss,val_uid_loss" >> {backup_writefile}')
##### end jason #####
while lr > args.min_lr and epoch_itr.next_epoch_idx <= max_epoch:
# train for one epoch
valid_losses, should_stop, train_stats, valid_stats = train(args, trainer, task, epoch_itr)
print("hello", valid_stats, train_stats)
##### begin jason #####
if train_stats and valid_stats:
updates_list.append(train_stats['num_updates'])
train_loss_list.append(train_stats['loss'])
train_ppl_list.append(train_stats['ppl'])
val_loss_list.append(valid_stats['loss'])
val_ppl_list.append(valid_stats['ppl'])
if 'uid_loss' not in train_stats:
train_stats['uid_loss'] = -1
valid_stats['uid_loss'] = -1
train_uid_loss_list.append(train_stats['uid_loss'])
val_uid_loss_list.append(valid_stats['uid_loss'])
log_line = f"{train_stats['num_updates']},{train_stats['loss']},{train_stats['ppl']},{valid_stats['loss']},{valid_stats['ppl']},{train_stats['uid_loss']},{valid_stats['uid_loss']}"
log_writer.write(f"{log_line}\n")
os.system(f'echo "{log_line}" >> {backup_writefile}')
best_val_loss = min(val_loss_list)
best_val_loss_idx = val_loss_list.index(best_val_loss)
updates_to_best_val_loss = updates_list[best_val_loss_idx]
train_loss_at_best_val_loss = train_loss_list[best_val_loss_idx]
jasons_vis.plot_jasons_lineplot(
x_list = updates_list,
y_list_list = [train_loss_list, val_loss_list, train_uid_loss_list, val_uid_loss_list],
y_labels_list = ['train', 'dev', 'train uid', 'dev uid'],
x_ax_label = "Updates",
y_ax_label = "Loss",
title = f"dev_l={best_val_loss} updates={updates_to_best_val_loss} train_l={train_loss_at_best_val_loss}",
output_png_path = os.path.join(args.jason_log_dir, f"{args.jason_log_dir.split('/')[-1]}_loss.png"),
)
jasons_vis.plot_jasons_lineplot(
x_list = updates_list,
y_list_list = [train_ppl_list, val_ppl_list],
y_labels_list = ['train', 'dev'],
x_ax_label = "Updates",
y_ax_label = "Perplexity",
title = f" best_val_ppl={best_val_loss} " + args.jason_log_dir[:20],
output_png_path = os.path.join(args.jason_log_dir, f"{args.jason_log_dir.split('/')[-1]}_perplexity.png"),
)
##### end jason #####
if should_stop:
break
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
epoch_itr = trainer.get_train_iterator(
epoch_itr.next_epoch_idx,
# sharded data: get train iterator for next epoch
load_dataset=task.has_sharded_data("train"),
# don't cache epoch iterators for sharded datasets
disable_iterator_cache=task.has_sharded_data("train"),
)
train_meter.stop()
logger.info("done training in {:.1f} seconds".format(train_meter.sum))
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
## 单机多卡和多机多卡训练都会调用这个函数
## 此函数中调用init_process_group函数，
## 此时还没有load数据，因此应该就没有了之前版本多机训练时因为load数据速度不同导致的超时问题
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args): ## 判断当前GPU是否是master GPU（args.distributed_rank = 0）
checkpoint_utils.verify_checkpoint_directory(args.save_dir) ## 确认checkpoint的目标存储路径
# Print args
print(args)
# Setup task, e.g., translation, language modeling, etc.
## 创建对应的TranslationTask类，读入两个dictionary: self.src_dict, self.tgt_dict, 并确定是left paddig or right padding
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
# 用于验证的开发集, 每个集合的名字为valid_sub_split。load之后，根据valid_sub_split的名字存放在task.datasets中
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args) ## 搭建神经网络模型, 翻译即使用TransformerModel类, 继承自FairseqEncoderDecoderModel
criterion = task.build_criterion(args) ## 搭建loss函数, 此处即使用LabelSmoothedCrossEntropyCriterion
print(model)
print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
##print the number of parameters of each matrix
#for name, param in model.named_parameters(recurse=True):
# print (name, param.numel())
#exit(0)
# Build trainer
# 如果distributed_world_size > 1, 则会对model和criterion使用models.DistributedFairseqModel进行wrap
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer) ## generate data iterator, epoch_itr
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
while (
lr > args.min_lr
and (epoch_itr.epoch < max_epoch or (epoch_itr.epoch == max_epoch
and epoch_itr._next_epoch_itr is not None))
and trainer.get_num_updates() < max_update
):
# train for one epoch
train(args, trainer, task, epoch_itr)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr, valid_subsets)
else:
valid_losses = [None]
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
##每个epoch都新建一个epoch data iterator来遍历所有的训练数据
reload_dataset = ':' in getattr(args, 'data', '')
# sharded data: get train iterator for next epoch
epoch_itr = trainer.get_train_iterator(epoch_itr.epoch, load_dataset=reload_dataset)
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
metrics.reset()
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
logger.info(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=1)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
logger.info(model)
logger.info('model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
logger.info('num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# (optionally) Configure quantization
if args.quantization_config_path is not None:
quantizer = quantization_utils.Quantizer(
config_path=args.quantization_config_path,
max_epoch=args.max_epoch,
max_update=args.max_update,
)
else:
quantizer = None
# Build trainer
if args.model_parallel_size == 1:
trainer = Trainer(args, task, model, criterion, quantizer)
else:
trainer = MegatronTrainer(args, task, model, criterion)
logger.info('training on {} GPUs'.format(args.distributed_world_size))
logger.info(
'max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = meters.StopwatchMeter()
train_meter.start()
while (lr > args.min_lr and epoch_itr.next_epoch_idx <= max_epoch):
# train for one epoch
valid_losses = train(args, trainer, task, epoch_itr, max_update)
if should_stop_early(
args,
valid_losses[0]) or trainer.get_num_updates() >= max_update:
break
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
epoch_itr = trainer.get_train_iterator(
epoch_itr.next_epoch_idx,
# sharded data: get train iterator for next epoch
load_dataset=(os.pathsep in getattr(args, 'data', '')),
)
train_meter.stop()
logger.info('done training in {:.1f} seconds'.format(train_meter.sum))
def main(
args,
init_distributed=False,
after_distributed_init_fn: Optional[Callable[[argparse.Namespace],
argparse.Namespace]] = None,
):
utils.import_user_module(args)
assert (
args.max_tokens is not None or args.max_sentences is not None
), "Must specify batch size either with --max-tokens or --max-sentences"
metrics.reset()
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu and not getattr(
args, "tpu", False):
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
utils.set_torch_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if after_distributed_init_fn:
args = after_distributed_init_fn(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
logger.info(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(","):
task.load_dataset(valid_sub_split, combine=False, epoch=1)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
logger.info(model)
logger.info("model {}, criterion {}".format(args.arch,
criterion.__class__.__name__))
logger.info("num. model params: {} (num. trained: {})".format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# (optionally) Configure quantization
if args.quantization_config_path is not None:
quantizer = quantization_utils.Quantizer(
config_path=args.quantization_config_path,
max_epoch=args.max_epoch,
max_update=args.max_update,
)
else:
quantizer = None
# Build trainer
if args.model_parallel_size == 1:
trainer = Trainer(args, task, model, criterion, quantizer)
else:
trainer = MegatronTrainer(args, task, model, criterion)
logger.info("training on {} devices (GPUs/TPUs)".format(
args.distributed_world_size))
logger.info("training on {} devices (GPUs/TPUs)".format(
args.distributed_world_size))
logger.info(
"max tokens per GPU = {} and max sentences per GPU = {}".format(
args.max_tokens, args.max_sentences))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
if args.tpu:
import torch_xla.core.xla_model as xm
xm.rendezvous("load_checkpoint") # wait for all workers
xm.mark_step()
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
lr = trainer.get_lr()
train_meter = meters.StopwatchMeter()
train_meter.start()
experiment_path = args.mhr_experiment # path for experiment configuration
total_samples = 0
restore = {
'enc_self_attn': None,
'dec_self_attn': None,
'dec_enc_attn': None
}
last_epoch_num = {
'enc_self_attn': 0,
'dec_self_attn': 0,
'dec_enc_attn': 0
}
while lr > args.min_lr and epoch_itr.next_epoch_idx <= max_epoch:
# train for one epoch
valid_losses, should_stop, total_samples_temp, restore, last_epoch_num = train(
args,
trainer,
task,
epoch_itr,
model,
experiment_path,
total_samples=total_samples,
restore=restore,
last_epoch_num=last_epoch_num)
total_samples = total_samples_temp
if should_stop:
break
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
epoch_itr = trainer.get_train_iterator(
epoch_itr.next_epoch_idx,
# sharded data: get train iterator for next epoch
load_dataset=(os.pathsep in getattr(args, "data", "")),
)
train_meter.stop()
logger.info("done training in {:.1f} seconds".format(train_meter.sum))
def main(cfg: DictConfig) -> None:
if isinstance(cfg, argparse.Namespace):
cfg = convert_namespace_to_omegaconf(cfg)
utils.import_user_module(cfg.common)
assert (
cfg.dataset.max_tokens is not None
or cfg.dataset.batch_size is not None
), "Must specify batch size either with --max-tokens or --batch-size"
metrics.reset()
np.random.seed(cfg.common.seed)
utils.set_torch_seed(cfg.common.seed)
if distributed_utils.is_master(cfg.distributed_training):
checkpoint_utils.verify_checkpoint_directory(cfg.checkpoint.save_dir)
# Print args
logger.info(cfg)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(cfg.task)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in cfg.dataset.valid_subset.split(","):
task.load_dataset(valid_sub_split, combine=False, epoch=1)
assert cfg.criterion, "Please specify criterion to train a model"
# Build model and criterion
model = task.build_model(cfg.model)
criterion = task.build_criterion(cfg.criterion)
logger.info(model)
logger.info("task: {}".format(task.__class__.__name__))
logger.info("model: {}".format(model.__class__.__name__))
logger.info("criterion: {})".format(criterion.__class__.__name__))
logger.info("num. model params: {} (num. trained: {})".format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# (optionally) Configure quantization
if cfg.common.quantization_config_path is not None:
quantizer = quantization_utils.Quantizer(
config_path=cfg.common.quantization_config_path,
max_epoch=cfg.optimization.max_epoch,
max_update=cfg.optimization.max_update,
)
else:
quantizer = None
# Build trainer
if cfg.common.model_parallel_size == 1:
trainer = Trainer(cfg, task, model, criterion, quantizer)
else:
trainer = MegatronTrainer(cfg, task, model, criterion)
logger.info("training on {} devices (GPUs/TPUs)".format(
cfg.distributed_training.distributed_world_size))
logger.info("max tokens per GPU = {} and batch size per GPU = {}".format(
cfg.dataset.max_tokens,
cfg.dataset.batch_size,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(
cfg.checkpoint,
trainer,
# don't cache epoch iterators for sharded datasets
disable_iterator_cache=task.has_sharded_data("train"),
)
max_epoch = cfg.optimization.max_epoch or math.inf
lr = trainer.get_lr()
train_meter = meters.StopwatchMeter()
train_meter.start()
while lr > cfg.optimization.min_lr and epoch_itr.next_epoch_idx <= max_epoch:
# train for one epoch
valid_losses, should_stop = train(cfg, trainer, task, epoch_itr)
if should_stop:
break
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
epoch_itr = trainer.get_train_iterator(
epoch_itr.next_epoch_idx,
# sharded data: get train iterator for next epoch
load_dataset=task.has_sharded_data("train"),
# don't cache epoch iterators for sharded datasets
disable_iterator_cache=task.has_sharded_data("train"),
)
train_meter.stop()
logger.info("done training in {:.1f} seconds".format(train_meter.sum))
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
print(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
adv_criterion = task.build_adversarial_criterion(args)
adv = task.build_adversary(args, model)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = AdversarialTrainer(args, task, model, criterion, adv_criterion,
adv)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr, filtered_maxpos_indices = checkpoint_utils.load_checkpoint(
args, trainer)
# pretrain data actor
if args.pretrain_data_actor and args.data_actor == 'lan' and args.data_actor_step_update:
trainer.pretrain_data_actor()
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
if args.eval_bleu:
gen_args = copy.deepcopy(args)
gen_args.sample = False
gen_args.beam = 5
gen_args.batch_size = 32
generator = task.build_generator(gen_args)
args.maximize_best_checkpoint_metric = True
else:
generator = None
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
epoch_itr = train(args, trainer, task, epoch_itr, generator,
filtered_maxpos_indices)
#trainer.update_language_sampler(args)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets, generator)
else:
valid_losses = [None]
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr,
valid_losses[0])
if ':' in getattr(args, 'data', ''):
# sharded data: get train iterator for next epoch
epoch_itr = trainer.get_train_iterator(epoch_itr.epoch)[0]
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
for idx in sorted(trainer.idx_to_dev_grad_dotprod.keys()):
print(idx)
str_dotprod = [str(i) for i in trainer.idx_to_dev_grad_dotprod[idx]]
print(" ".join(str_dotprod))
def main(args):
utils.import_user_module(args)
assert (
args.max_tokens is not None or args.max_sentences is not None
), "Must specify batch size either with --max-tokens or --max-sentences"
metrics.reset()
np.random.seed(args.seed)
utils.set_torch_seed(args.seed)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
logger.info(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(","):
task.load_dataset(valid_sub_split, combine=False, epoch=1)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
logger.info(model)
logger.info("task: {} ({})".format(args.task, task.__class__.__name__))
logger.info("model: {} ({})".format(args.arch, model.__class__.__name__))
logger.info("criterion: {} ({})".format(args.criterion,
criterion.__class__.__name__))
logger.info("num. model params: {} (num. trained: {})".format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# (optionally) Configure quantization
if args.quantization_config_path is not None:
quantizer = quantization_utils.Quantizer(
config_path=args.quantization_config_path,
max_epoch=args.max_epoch,
max_update=args.max_update,
)
else:
quantizer = None
# Build trainer
if args.model_parallel_size == 1:
trainer = Trainer(args, task, model, criterion, quantizer)
else:
trainer = MegatronTrainer(args, task, model, criterion)
logger.info("training on {} devices (GPUs/TPUs)".format(
args.distributed_world_size))
logger.info(
"max tokens per GPU = {} and max sentences per GPU = {}".format(
args.max_tokens, args.max_sentences))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
lr = trainer.get_lr()
train_meter = meters.StopwatchMeter()
train_meter.start()
while lr > args.min_lr and epoch_itr.next_epoch_idx <= max_epoch:
# train for one epoch
valid_losses, should_stop = train(args, trainer, task, epoch_itr)
if should_stop:
break
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
epoch_itr = trainer.get_train_iterator(
epoch_itr.next_epoch_idx,
# sharded data: get train iterator for next epoch
load_dataset=task.has_sharded_data("train"),
)
train_meter.stop()
logger.info("done training in {:.1f} seconds".format(train_meter.sum))
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_init_hvd(args)
# Print args
print(args)
# if not HAS_NSML:
# args.data[0] = args.data[0].replace("/train", "")
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
if args.train_decoder_only:
for name, param in model.named_parameters():
if "decoder" not in name:
param.requires_grad_(False)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Setup session
if HAS_WANDB and distributed_utils.is_master(args):
wandb.init(project="cmlm", config=args)
wandb.watch(model)
# Load pre-trained model
data_token = args.data[0].split("/")[-1]
if "bert" in args.arch:
pretrained_path = "{}/train/pretrained_models/maskPredict_{}/checkpoint_best.pt".format(
DATASET_PATH,
data_token.split(".")[-1].replace("-", "_"))
if not HAS_NSML:
pretrained_path = pretrained_path.replace("/train", "")
print("| loading", pretrained_path)
state = checkpoint_utils.load_checkpoint_to_cpu(pretrained_path)
model.load_state_dict(state["model"], strict=True)
baseline_model = task.build_model(args)
baseline_model.load_state_dict(state["model"], strict=True)
if torch.cuda.is_available():
baseline_model.cuda()
task.set_baseline_model(baseline_model)
if not args.masking and HAS_NSML:
def nsml_bind(model):
def save(dir_path):
state = {
'model': model.state_dict(),
}
torch.save(state, os.path.join(dir_path, 'best.pt'))
def load(dir_path):
state = torch.load(os.path.join(dir_path, 'best.pt'),
map_location="cpu")
model.load_state_dict(state['model'], strict=False)
model.cuda()
print('model loaded!')
nsml.bind(save=save, load=load)
nsml_bind(model)
if args.load:
print("loading model from session", args.load)
if args.load.startswith("nsml://"):
session = args.load.replace("nsml://", "")
if ".pt" in session:
session = session.replace(".pt", "")
session, checkpoint_name = session.rsplit("/", 1)
else:
checkpoint_name = "best"
if "-" in checkpoint_name:
start, end = checkpoint_name.replace("epoch", "").split("-")
checkpoints = [
"epoch{}".format(i) for i in range(int(start),
int(end) + 1)
]
print("| checkpoint average:", checkpoints)
state_dict = None
def load(dir_path):
nonlocal state_dict, checkpoints
state = torch.load(os.path.join(dir_path, 'best.pt'))
model_state = state["model"]
for k in model_state:
model_state[k] = model_state[k] / float(len(checkpoints))
if state_dict is None:
state_dict = model_state
else:
for k in state_dict:
state_dict[k] += model_state[k]
print("checkpoint loaded")
for checkpoint_name in checkpoints:
nsml.load(checkpoint_name, load_fn=load, session=session)
model.load_state_dict(state_dict)
else:
def load(dir_path):
state = torch.load(os.path.join(dir_path, 'best.pt'))
state_dict = state["model"]
model.load_state_dict(state_dict)
print("loaded")
nsml.load(checkpoint_name, load_fn=load, session=session)
# Prepare for decoder wise training
if args.decoder_wise_training:
print("| Decoder wise training, start refinement step 0")
progressive_training_step = 0
assert args.ddp_backend == "c10d"
else:
progressive_training_step = None
# Build trainer
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
if hasattr(args, "progressive") and args.progressive:
for i in range(args.refinetot if not getattr(args, "pnet", False) else
args.refinetot - 1):
print("validating for refine step", i)
validate(args,
trainer,
task,
epoch_itr,
valid_subsets,
force_refine_step=i)
print("---")
validate(args, trainer, task, epoch_itr, valid_subsets)
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args,
trainer,
task,
epoch_itr,
force_refine_step=progressive_training_step)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(
args,
trainer,
task,
epoch_itr,
valid_subsets,
force_refine_step=progressive_training_step)
else:
valid_losses = [None]
if args.decoder_wise_training:
progressive_training_step = update_num_to_refine_step(
trainer.get_num_updates())
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
if HAS_NSML:
if distributed_utils.is_master(args):
print("nsml save for epoch", epoch_itr.epoch)
nsml.save("epoch{}".format(epoch_itr.epoch))
else:
torch.save({"model": trainer.get_model().state_dict()},
"/tmp/epoch{}.pt".format(epoch_itr.epoch))
if HAS_WANDB:
wandb.save("/tmp/epoch{}.pt".format(epoch_itr.epoch))
# checkpoint_utils.save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
if ':' in getattr(args, 'data', ''):
# sharded data: get train iterator for next epoch
epoch_itr = trainer.get_train_iterator(epoch_itr.epoch)
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
print(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
else:
valid_losses = [None]
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr,
valid_losses[0])
reload_dataset = ':' in getattr(args, 'data', '')
# sharded data: get train iterator for next epoch
epoch_itr = trainer.get_train_iterator(epoch_itr.epoch,
load_dataset=reload_dataset)
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
# Print args
print(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
task.load_dataset(args.train_subset, combine=True, epoch=0)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=True, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=args.required_batch_size_multiple,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
num_workers=args.num_workers,
)
# Load the latest checkpoint if one is available
checkpoint_utils.load_checkpoint(args, trainer, epoch_itr, max_positions,
task)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr,
valid_losses[0])
epoch_itr = checkpoint_utils.reload_train(args, epoch_itr,
max_positions, task)
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
