def test(opt):
# 数据
dataloader = get_dataloader(opt)
_data = dataloader.dataset._data
word2ix, ix2word = _data['word2ix'], _data['ix2word']
sos = word2ix.get(_data.get('sos'))
eos = word2ix.get(_data.get('eos'))
unknown = word2ix.get(_data.get('unknown'))
voc_length = len(word2ix)
#定义模型
encoder = EncoderRNN(opt, voc_length)
decoder = LuongAttnDecoderRNN(opt, voc_length)
#加载模型
if opt.model_ckpt == None:
raise ValueError('model_ckpt is None.')
return False
checkpoint = torch.load(opt.model_ckpt, map_location=lambda s, l: s)
encoder.load_state_dict(checkpoint['en'])
decoder.load_state_dict(checkpoint['de'])
with torch.no_grad():
#切换模式
encoder = encoder.to(opt.device)
decoder = decoder.to(opt.device)
encoder.eval()
decoder.eval()
#定义seracher
searcher = GreedySearchDecoder(encoder, decoder)
return searcher, sos, eos, unknown, word2ix, ix2word
def eval(**kwargs):
opt = Config()
for k, v in kwargs.items(): #设置参数
setattr(opt, k, v)
# 数据
dataloader = get_dataloader(opt)
_data = dataloader.dataset._data
word2ix, ix2word = _data['word2ix'], _data['ix2word']
sos = word2ix.get(_data.get('sos'))
eos = word2ix.get(_data.get('eos'))
unknown = word2ix.get(_data.get('unknown'))
voc_length = len(word2ix)
#定义模型
encoder = EncoderRNN(opt, voc_length)
decoder = LuongAttnDecoderRNN(opt, voc_length)
#加载模型
if opt.model_ckpt == None:
raise ValueError('model_ckpt is None.')
return False
checkpoint = torch.load(opt.model_ckpt, map_location=lambda s, l: s)
encoder.load_state_dict(checkpoint['en'])
decoder.load_state_dict(checkpoint['de'])
with torch.no_grad():
#切换模式
encoder = encoder.to(opt.device)
decoder = decoder.to(opt.device)
encoder.eval()
decoder.eval()
#定义seracher
searcher = GreedySearchDecoder(encoder, decoder)
while (1):
input_sentence = input('> ')
if input_sentence == 'q' or input_sentence == 'quit': break
cop = re.compile("[^\u4e00-\u9fa5^a-z^A-Z^0-9]") #分词处理正则
input_seq = jieba.lcut(cop.sub("", input_sentence)) #分词序列
input_seq = input_seq[:opt.max_input_length] + ['</EOS>']
input_seq = [word2ix.get(word, unknown) for word in input_seq]
tokens = generate(input_seq, searcher, sos, eos, opt)
output_words = ''.join([ix2word[token.item()] for token in tokens])
print('BOT: ', output_words)
def train(**kwargs):
opt = Config()
for k, v in kwargs.items(): #设置参数
setattr(opt, k, v)
# 数据
dataloader = get_dataloader(opt)
_data = dataloader.dataset._data
word2ix = _data['word2ix']
sos = word2ix.get(_data.get('sos'))
voc_length = len(word2ix)
#定义模型
encoder = EncoderRNN(opt, voc_length)
decoder = LuongAttnDecoderRNN(opt, voc_length)
#加载断点,从上次结束地方开始
if opt.model_ckpt:
checkpoint = torch.load(opt.model_ckpt)
encoder.load_state_dict(checkpoint['en'])
decoder.load_state_dict(checkpoint['de'])
#切换模式
encoder = encoder.to(opt.device)
decoder = decoder.to(opt.device)
encoder.train()
decoder.train()
#定义优化器(注意与encoder.to(device)前后不要反)
encoder_optimizer = torch.optim.Adam(encoder.parameters(),
lr=opt.learning_rate)
decoder_optimizer = torch.optim.Adam(decoder.parameters(),
lr=opt.learning_rate *
opt.decoder_learning_ratio)
if opt.model_ckpt:
encoder_optimizer.load_state_dict(checkpoint['en_opt'])
decoder_optimizer.load_state_dict(checkpoint['de_opt'])
#定义打印loss的变量
print_loss = 0
for epoch in range(opt.epoch):
for ii, data in enumerate(dataloader):
#取一个batch训练
loss = train_by_batch(sos, opt, data, encoder_optimizer,
decoder_optimizer, encoder, decoder)
print_loss += loss
#打印损失
if ii % opt.print_every == 0:
print_loss_avg = print_loss / opt.print_every
print(
"Epoch: {}; Epoch Percent complete: {:.1f}%; Average loss: {:.4f}"
.format(epoch, epoch / opt.epoch * 100, print_loss_avg))
print_loss = 0
# 保存checkpoint
if epoch % opt.save_every == 0:
checkpoint_path = '{prefix}_{time}'.format(
prefix=opt.prefix, time=time.strftime('%m%d_%H%M'))
torch.save(
{
'en': encoder.state_dict(),
'de': decoder.state_dict(),
'en_opt': encoder_optimizer.state_dict(),
'de_opt': decoder_optimizer.state_dict(),
}, checkpoint_path)
print('param num = {}'.format(para_num))
if torch.cuda.is_available():
model.cuda()
# prepare optimizer,lr scheduler,loss function
optimizer = torch.optim.AdamW(model.parameters(),
lr=opt.lr,
weight_decay=1e-2)
scheduler_S = optim.lr_scheduler.CosineAnnealingLR(optimizer, EPOCHES,
1e-6)
loss_fn = ContourCE()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
# prepare data
train_loader, val_loader = get_dataloader()
# prepare metric
best_iou = None
# prepare weight storage path
# save_path = './weight/denseunet_8_images.pth'
# load_path = './weight/densenet9.pth'
# load_path = './weight/densenet9_v3.pth'
load_path = './weight/densenet9_v4_batch3.pth'
if os.path.exists(load_path):
model.load_state_dict(torch.load(load_path))
model.eval()
print('load from {}'.format(load_path))
# 验证集上的估计
val_loss, val_iou = evaluate_avg(model,
val_loader,