a = F.softmax(u / self.softmax_T, dim=1)
d = torch.bmm(u2, a.unsqueeze(2)).squeeze(2)
return u, d.unsqueeze(1)
def get_log_likelihood(self, _log_p, pi):
""" args:
_log_p: (batch, city_t, city_t)
pi: (batch, city_t), predicted tour
return: (batch)
"""
log_p = torch.gather(input=_log_p, dim=2, index=pi[:, :, None])
return torch.sum(log_p.squeeze(-1), 1)
if __name__ == '__main__':
cfg = load_pkl(pkl_parser().path)
model = PtrNet1(cfg)
inputs = torch.randn(3, 20, 2)
pi, ll = model(inputs, device='cpu')
print('pi:', pi.size(), pi)
print('log_likelihood:', ll.size(), ll)
cnt = 0
for i, k in model.state_dict().items():
print(i, k.size(), torch.numel(k))
cnt += torch.numel(k)
print('total parameters:', cnt)
# ll.mean().backward()
# print(model.W_q.weight.grad)
grads, _ = tf.clip_by_global_norm(grads, 1.0)
optimizer.apply_gradients(zip(grads, model.trainable_weights))# optimizer.step
ave_loss.update_state(loss)
ave_L.update_state(L_mean)
if t%(cfg.batch_steps*0.1) == 0:
print('epoch%d, %d/%dsamples: loss %1.2f, average L %1.2f, average b %1.2f\n'%(
epoch, t*cfg.batch, cfg.n_samples, ave_loss.result().numpy(), ave_L.result().numpy(), tf.reduce_mean(b)))
baseline.epoch_callback(model, epoch)
model.decode_type = 'sampling'
model.save_weights('%s%s_epoch%s.h5'%(cfg.weight_dir, cfg.task, epoch), save_format = 'h5')
if cfg.islogger:
if log_path is None:
log_path = '%s%s_%s.csv'%(cfg.log_dir, cfg.task, cfg.dump_date)#cfg.log_dir = ./Csv/
with open(log_path, 'w') as f:
f.write('time,epoch,loss,average length\n')
with open(log_path, 'a') as f:
t2 = time()
f.write('%dmin%dsec,%d,%1.2f,%1.2f\n'%((t2-t1)//60, (t2-t1)%60, epoch, ave_loss.result().numpy(), ave_L.result().numpy()))
ave_loss.reset_states()
ave_L.reset_states()
if __name__ == '__main__':
cfg = load_pkl(file_parser().path)
train(cfg)
print('update min_val_avg_loss: ', min_val_avg_loss)
weight_path = '%s%s_epoch%s.pt' % (cfg.weight_dir, cfg.dump_date,
epoch)
torch.save(model.state_dict(), weight_path)
print(f'generate {weight_path}')
if cfg.islogger:
if epoch == 1:
val_csv_path = '%s%s_val.csv' % (cfg.log_dir, cfg.dump_date
) #cfg.log_dir = ./Csv/
print(f'generate {val_csv_path}')
with open(val_csv_path, 'w') as f:
f.write('time,epoch,val avg loss\n')
with open(val_csv_path, 'a') as f:
f.write('%dmin%dsec,%d,%1.3f\n' %
((t6 - t5) // 60,
(t6 - t5) % 60, epoch, val_avg_loss / t))
else: # val_avg_loss/t >= min_val_avg_loss:
cnt += 1
print(f'current cnt: {cnt}/7')
if cnt >= 7:
print('early stopping')
break
if __name__ == '__main__':
cfg = load_pkl(train_parser().path)
train(cfg)