def train(epoch):
data_loader = get_dataloader(train=True)
bar = tqdm(data_loader, total=len(data_loader))
for idx, (input, target, input_len, target_len) in enumerate(bar):
input = input.to(config.device)
target = target.to(config.device)
input_len = input_len.to(config.device)
optimizer.zero_grad()
decoder_outputs = model(input, input_len,
target) #[batch_Size,max_len,vocab_size]
loss = F.nll_loss(decoder_outputs.view(-1, len(config.target_ws)),
target.view(-1),
ignore_index=config.input_ws.PAD)
loss.backward()
optimizer.step()
loss_list.append(loss.item())
bar.set_description("epoch:{} idx:{} loss:{:.6f}".format(
epoch, idx, np.mean(loss_list)))
if idx % 100 == 0:
torch.save(model.state_dict(), "../chatbot/models/model.pkl")
torch.save(optimizer.state_dict(),
"../chatbot/models/optimizer.pkl")
pickle.dump(loss_list, open("../chatbot/models/loss_list.pkl",
"wb"))
def test_encoder():
train_dataloader, test_dataloader = get_dataloader()
encoder = Encoder()
print(encoder)
for feature, target, feature_length, target_length in train_dataloader:
output, hidden = encoder(feature, feature_length)
print(output.size())
print(hidden.size())
break
def test_chatbot_dataloader():
"""测试dataloader"""
train_dataloader, test_dataloader = get_dataloader()
for input, target, input_length, target_length in train_dataloader:
print(input)
print(target)
print(input_length)
print(target_length)
break
def test_dataloader():
loader = get_dataloader()
for idx, (input, target, input_len, target_len) in enumerate(loader):
print(idx)
print(input)
print(target)
print(input_len)
print(target_len)
break
def test_decoder():
train_dataloader, test_dataloader = get_dataloader()
encoder = Encoder()
decoder = Decoder()
print(encoder)
print(decoder)
for feature, target, feature_length, target_length in train_dataloader:
output, encoder_hidden = encoder(feature, feature_length)
result = decoder(encoder_hidden, target)
print(result.size())
break
def eval():
model = Seq2Seq().to(config.device)
model.eval()
model.load_state_dict(torch.load("./models/model.pkl"))
loss_list = []
data_loader = get_dataloader(train=False)
bar = tqdm(data_loader, total=len(data_loader), desc="当前进行评估")
with torch.no_grad():
for idx, (input, target, input_len, target_len) in enumerate(bar):
input = input.to(config.device)
target = target.to(config.device)
input_len = input_len.to(config.device)
decoder_outputs, predict_result = model.evaluate(
input, input_len) #[batch_Size,max_len,vocab_size]
loss = F.nll_loss(decoder_outputs.view(-1, len(config.target_ws)),
target.view(-1),
ignore_index=config.input_ws.PAD)
loss_list.append(loss.item())
bar.set_description("idx:{} loss:{:.6f}".format(
idx, np.mean(loss_list)))
print("当前的平均损失为：", np.mean(loss_list))
from lib import cut
from config import by_char
# 使用gpu训练
s2s = Seq2Seq().to(config.device)
optimizer = torch.optim.Adam(s2s.parameters())
# 模型加载
if os.path.exists(config.chatbot_model_path) and os.path.exists(
config.chatbot_optimizer_path):
s2s.load_state_dict(
torch.load(config.chatbot_model_path, map_location=config.device))
optimizer.load_state_dict(
torch.load(config.chatbot_optimizer_path, map_location=config.device))
train_dataloader, test_dataloader = get_dataloader()
def train():
for i in range(config.EPOCHS):
with tqdm(total=len(train_dataloader)) as t:
for index, (feature, target, feature_length,
target_length) in enumerate(train_dataloader):
t.set_description("Epoch %i/%i" %
(i + 1, config.EPOCHS)) # 设置描述
optimizer.zero_grad() # 梯度置为0
feature = feature.to(config.device)
target = target.to(config.device)
feature_length = feature_length.to(config.device)
target_length = target_length.to(config.device)
y_predict = s2s(feature, feature_length,