def load_model(filepath):
"""
Load a MusicTransformer from a saved pytorch state_dict and hparams. The input filepath should point to a .pt
file in which has been saved a dictionary containing the model state dict and hparams, ex:
torch.save(filepath, {
"state_dict": MusicTransformer.state_dict(),
"hparams": hparams (dict)
})
Args:
filepath (str): path to single .pt file containing the dictionary as described above
Returns:
the loaded MusicTransformer model
"""
from model import MusicTransformer
from hparams import hparams
file = torch.load(filepath)
if "hparams" not in file:
file["hparams"] = hparams
model = MusicTransformer(**file["hparams"]).to(device)
model.load_state_dict(file["state_dict"])
model.eval()
return model
class MusicTransformerTrainer:
"""
As the transformer is a large model and takes a while to train on a GPU, or even a TPU, I wrote this Trainer
class to make it easier to load and save checkpoints with the model. The way I've designed it instantiates the
model, optimizer, and scheduler within the class itself, as there are some problems with passing them in. But,
to get these objects back just call:
trainer.model
trainer.optimizer
trainer.scheduler
This class also tracks the cumulative losses while training, which you can get back with:
trainer.train_losses
trainer.val_losses
as lists of floats
To save a checkpoint, call trainer.save()
To load a checkpoint, call trainer.load( (optional) ckpt_path)
"""
def __init__(self, hparams_, datapath, batch_size, warmup_steps=4000,
ckpt_path="music_transformer_ckpt.pt", load_from_checkpoint=False):
"""
Args:
hparams_: hyperparameters of the model
datapath: path to the data to train on
batch_size: batch size to batch the data
warmup_steps: number of warmup steps for transformer learning rate schedule
ckpt_path: path at which to save checkpoints while training; MUST end in .pt or .pth
load_from_checkpoint (bool, optional): if true, on instantiating the trainer, this will load a previously
saved checkpoint at ckpt_path
"""
# get the data
self.datapath = datapath
self.batch_size = batch_size
data = torch.load(datapath).long().to(device)
# max absolute position must be able to acount for the largest sequence in the data
if hparams_["max_abs_position"] > 0:
hparams_["max_abs_position"] = max(hparams_["max_abs_position"], data.shape[-1])
# train / validation split: 80 / 20
train_len = round(data.shape[0] * 0.8)
train_data = data[:train_len]
val_data = data[train_len:]
print(f"There are {data.shape[0]} samples in the data, {len(train_data)} training samples and {len(val_data)} "
"validation samples")
# datasets and dataloaders: split data into first (n-1) and last (n-1) tokens
self.train_ds = TensorDataset(train_data[:, :-1], train_data[:, 1:])
self.train_dl = DataLoader(dataset=self.train_ds, batch_size=batch_size, shuffle=True)
self.val_ds = TensorDataset(val_data[:, :-1], val_data[:, 1:])
self.val_dl = DataLoader(dataset=self.val_ds, batch_size=batch_size, shuffle=True)
# create model
self.model = MusicTransformer(**hparams_).to(device)
self.hparams = hparams_
# setup training
self.warmup_steps = warmup_steps
self.optimizer = optim.Adam(self.model.parameters(), lr=1.0, betas=(0.9, 0.98))
self.scheduler = optim.lr_scheduler.LambdaLR(
self.optimizer,
lambda x: transformer_lr_schedule(self.hparams['d_model'], x, self.warmup_steps)
)
# setup checkpointing / saving
self.ckpt_path = ckpt_path
self.train_losses = []
self.val_losses = []
# load checkpoint if necessesary
if load_from_checkpoint and os.path.isfile(self.ckpt_path):
self.load()
def save(self, ckpt_path=None):
"""
Saves a checkpoint at ckpt_path
Args:
ckpt_path (str, optional): if None, saves the checkpoint at the previously stored self.ckpt_path
else saves the checkpoints at the new passed-in path, and stores this new path at
the member variable self.ckpt_path
"""
if ckpt_path is not None:
self.ckpt_path = ckpt_path
ckpt = {
"model_state_dict": self.model.state_dict(),
"optimizer_state_dict": self.optimizer.state_dict(),
"scheduler_state_dict": self.scheduler.state_dict(),
"train_losses": self.train_losses,
"validation_losses": self.val_losses,
"warmup_steps": self.warmup_steps,
"hparams": self.hparams
}
torch.save(ckpt, self.ckpt_path)
return
def load(self, ckpt_path=None):
"""
Loads a checkpoint from ckpt_path
NOTE: OVERWRITES THE MODEL STATE DICT, OPTIMIZER STATE DICT, SCHEDULER STATE DICT, AND HISTORY OF LOSSES
Args:
ckpt_path (str, optional): if None, loads the checkpoint at the previously stored self.ckpt_path
else loads the checkpoints from the new passed-in path, and stores this new path
at the member variable self.ckpt_path
"""
if ckpt_path is not None:
self.ckpt_path = ckpt_path
ckpt = torch.load(self.ckpt_path)
del self.model, self.optimizer, self.scheduler
# create and load model
self.model = MusicTransformer(**ckpt["hparams"]).to(device)
self.hparams = ckpt["hparams"]
print("Loading the model...", end="")
print(self.model.load_state_dict(ckpt["model_state_dict"]))
# create and load load optimizer and scheduler
self.warmup_steps = ckpt["warmup_steps"]
self.optimizer = optim.Adam(self.model.parameters(), lr=1.0, betas=(0.9, 0.98))
self.optimizer.load_state_dict(ckpt["optimizer_state_dict"])
self.scheduler = optim.lr_scheduler.LambdaLR(
self.optimizer,
lambda x: transformer_lr_schedule(self.hparams['d_model'], x, self.warmup_steps)
)
self.scheduler.load_state_dict(ckpt["scheduler_state_dict"])
# load loss histories
self.train_losses = ckpt["train_losses"]
self.val_losses = ckpt["validation_losses"]
return
def fit(self, epochs):
"""
Training loop to fit the model to the data stored at the passed in datapath. If KeyboardInterrupt at anytime
during the training loop, and if progresss being printed, this method will save a checkpoint at the
passed-in ckpt_path
Args:
epochs: number of epochs to train for.
Returns:
history of training and validation losses for this training session
"""
print_interval = epochs // 10 + int(epochs < 10)
train_losses = []
val_losses = []
start = time.time()
print("Beginning training...")
try:
for epoch in range(epochs):
train_epoch_losses = []
val_epoch_losses = []
self.model.train()
for train_inp, train_tar in self.train_dl:
loss = train_step(self.model, self.optimizer, self.scheduler, train_inp, train_tar)
train_epoch_losses.append(loss)
self.model.eval()
for val_inp, val_tar in self.val_dl:
loss = val_step(self.model, val_inp, val_tar)
val_epoch_losses.append(loss)
# mean losses for the epoch
train_mean = sum(train_epoch_losses) / len(train_epoch_losses)
val_mean = sum(val_epoch_losses) / len(val_epoch_losses)
# store complete history of losses in member lists and relative history for this session in output lists
self.train_losses.append(train_mean)
train_losses.append(train_mean)
self.val_losses.append(val_mean)
val_losses.append(val_mean)
if ((epoch + 1) % print_interval) == 0:
print(f"Epoch {epoch + 1} Time taken {round(time.time() - start, 2)} seconds "
f"Train Loss {train_losses[-1]} Val Loss {val_losses[-1]}")
# print("Checkpointing...")
# self.save()
# print("Done")
start = time.time()
except KeyboardInterrupt:
pass
print("Checkpointing...")
self.save()
print("Done")
return train_losses, val_losses
