File size: 8,485 Bytes
c961996 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 |
import torch
from tqdm import tqdm
from data import LyricsCommentsDatasetPsuedo_fusion
from torch import utils, nn
from model_fusion import CommentGenerator_fusion
import transformers
import time
import statistics
import os
import random
import datasets
IS_LOAD = False
LOAD_EPOCH = 0
EPOCH = 50
BATCH_SIZE = 8
LOG_INTERVAL = 100
SAMPLE_INTERVAL = 1000
VALIDATION_INTERVAL = 2
LOG_FOLDER = "log/"
MODEL_FOLDER = "model/"
SAVE_INTERVAL = 2
EARLY_STOPPING_INTERVAL = 5
MODEL_NAME = "bart_fusion_full_256"
CHOICE_NUMBER = 2
DATASET_PATH = "/homes/yz007/multimodal-transformer/comment_generator/dataset_full_256.pkl"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
dataset = LyricsCommentsDatasetPsuedo_fusion(dataset_path=DATASET_PATH)
dataset_length = len(dataset)
train_dataset_length = int(dataset_length * 0.9)
valid_dataset_length = dataset_length - train_dataset_length
train_dataset, valid_dataset = utils.data.random_split(dataset,
[train_dataset_length,
valid_dataset_length],
generator=torch.Generator().manual_seed(42))
train_dataloader = utils.data.DataLoader(train_dataset,
batch_size=BATCH_SIZE,
shuffle=True)
# valid_dataloader = utils.data.DataLoader(valid_dataset,
# batch_size=32,
# shuffle=False)
model = CommentGenerator_fusion().cuda()
criterion = nn.CrossEntropyLoss()
# optimizer = transformers.Adafactor(filter(lambda p: p.requires_grad, model.parameters()),
# lr=6e-4,
# )
optimizer = transformers.Adafactor(model.parameters(), warmup_init=False, relative_step=False,
lr=6e-4,
)
if IS_LOAD:
model.load_state_dict(torch.load("/homes/yz007/multimodal-transformer/comment_generator/model/bart_fusion_positive_256_6e-4_epoch6.pt"))
optimizer.load_state_dict(torch.load("/homes/yz007/multimodal-transformer/comment_generator/model/bart_fusion_positive_256_6e-4_optim_epoch6.pt"))
loss_stat = list()
start_time = time.time()
start_time_local = time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime())
early_stop_token = [0.0, 0]
validation_loss_history = list()
model.train()
for epoch in range(1 + LOAD_EPOCH, EPOCH + 1 + LOAD_EPOCH):
for batch_index, [lyrics, comment, music_id] in enumerate(train_dataloader):
# pre-process data
input_sentences = lyrics
raw_labels = comment
output = model(input_sentences, music_id, raw_labels)
loss = output.loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_stat.append(loss.item())
# log
if batch_index and batch_index % LOG_INTERVAL == 0:
curr_time = time.time()
passed_time_all = curr_time - start_time
time_str = f"{int(passed_time_all / 60)}:{int(passed_time_all % 60)}"
log = f"{MODEL_NAME}\t" \
f"Time: {time_str}\t" \
f"Epoch {epoch}: {batch_index}/{int(len(train_dataloader.dataset) / BATCH_SIZE)}\t" \
f"Loss: {statistics.mean(loss_stat[-1 * LOG_INTERVAL * BATCH_SIZE:])}\t" \
f"Avg loss: {statistics.mean(loss_stat)}"
if __debug__:
print(log)
with open(os.path.join(LOG_FOLDER, MODEL_NAME + '_' + start_time_local + ".txt"), 'a+',
encoding='utf-8') as r:
r.write(log)
r.write("\n")
loss_stat = list()
if batch_index and batch_index % SAMPLE_INTERVAL == 0:
# make samples
model.eval()
samples_list = random.choices(valid_dataset, k=CHOICE_NUMBER)
sample_sentence, sample_label, music_ids = zip(*samples_list)
with torch.no_grad():
output_samples = model.generate(sample_sentence, music_ids)
for sample_index in range(CHOICE_NUMBER):
log = f"Lyrics: {sample_sentence[sample_index]}\n" \
f"Sample outputs: {output_samples[sample_index]}\n" \
f"Ground Truth: {sample_label[sample_index]}"
if __debug__:
print(log)
with open(os.path.join(LOG_FOLDER, MODEL_NAME + '_' + start_time_local + ".txt"), 'a+',
encoding='utf-8') as r:
r.write(log)
r.write("\n")
# validation loss
valid_dataloader = utils.data.DataLoader(valid_dataset,
batch_size=8,
shuffle=False)
valid_loss_stat = list()
for batch_index_valid, [lyrics_valid, comment_valid, music_id_valid] in enumerate(valid_dataloader):
with torch.no_grad():
output_valid = model(lyrics_valid, music_id_valid, comment_valid)
valid_loss = output_valid.loss.item()
valid_loss_stat.append(valid_loss)
if batch_index_valid > 15:
break
valid_loss_mean = statistics.mean(valid_loss_stat)
validation_loss_history.append(valid_loss_mean)
log = f"{MODEL_NAME}\t" \
f"Time: {time_str}\t" \
f"Epoch {epoch}: {batch_index}/{int(len(train_dataloader.dataset) / BATCH_SIZE)}\t" \
f"Validation Loss: {valid_loss_mean}\t"
if __debug__:
print(log)
with open(os.path.join(LOG_FOLDER, MODEL_NAME + '_' + start_time_local + ".txt"), 'a+',
encoding='utf-8') as r:
r.write(log)
r.write("\n")
# back to train
model.train()
if epoch and epoch % VALIDATION_INTERVAL == 0:
model.eval()
metrics = datasets.load_metric('rouge')
valid_dataloader = utils.data.DataLoader(valid_dataset,
batch_size=8,
shuffle=False)
for batch_index_valid, [lyrics_valid, comment_valid, music_id_valid] in enumerate(valid_dataloader):
with torch.no_grad():
output_samples = model.generate(lyrics_valid, music_id_valid)
metrics.add_batch(predictions=output_samples, references=comment_valid)
# control time.
if batch_index_valid > 10:
break
score = metrics.compute()
if __debug__:
print(str(score))
with open(os.path.join(LOG_FOLDER, MODEL_NAME + '_' + start_time_local + ".txt"), 'a+',
encoding='utf-8') as r:
r.write(str(score))
r.write("\n")
# save
if score['rouge1'].mid.recall > early_stop_token[0]:
early_stop_token = [score['rouge1'].mid.recall, epoch] # replace to the best
torch.save(model.state_dict(), os.path.join(MODEL_FOLDER, f"{MODEL_NAME}_best.pt"))
torch.save(optimizer.state_dict(),
os.path.join(MODEL_FOLDER, f"{MODEL_NAME}_optim_best.pt"))
# save
if epoch and epoch % SAVE_INTERVAL == 0:
torch.save(model.state_dict(), os.path.join(MODEL_FOLDER, f"{MODEL_NAME}_epoch{epoch}.pt"))
torch.save(optimizer.state_dict(),
os.path.join(MODEL_FOLDER, f"{MODEL_NAME}_optim_epoch{epoch}.pt"))
# early stopping
if len(validation_loss_history) > EARLY_STOPPING_INTERVAL:
if min(validation_loss_history[-2 * EARLY_STOPPING_INTERVAL:]) == validation_loss_history[-2 * EARLY_STOPPING_INTERVAL]:
print(f"Early Stopping. Best Score: {early_stop_token[0]} at Epoch {early_stop_token[1]}.")
break
if score['rouge1'].mid.recall <= early_stop_token[0] and epoch > (
early_stop_token[1] + EARLY_STOPPING_INTERVAL):
print(f"Early Stopping. Best Score: {early_stop_token[0]} at Epoch {early_stop_token[1]}.")
break
model.train()
print(f"Training Complete. Best Score: {early_stop_token[0]} at Epoch {early_stop_token[1]}.")
|