|
|
import os |
|
|
import sys |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
from datetime import datetime |
|
|
from torch.nn.utils import clip_grad_norm_ |
|
|
import torch.optim as optim |
|
|
from torch.utils.data import DataLoader |
|
|
import json |
|
|
|
|
|
|
|
|
dir_path = os.path.dirname(os.path.realpath(__file__)) |
|
|
sys.path.append(os.path.join(dir_path, '..')) |
|
|
|
|
|
from nets.model import Model |
|
|
from Actor.actor import Actor |
|
|
from dataloader import VRP_Dataset |
|
|
from google_solver.google_model import evaluate_google_model |
|
|
|
|
|
|
|
|
with open('params.json', 'r') as f: |
|
|
params = json.load(f) |
|
|
|
|
|
|
|
|
safe_data_dir = "/home/user/data" |
|
|
os.makedirs(safe_data_dir, exist_ok=True) |
|
|
|
|
|
|
|
|
with open(os.path.join(safe_data_dir, 'params_saved.json'), 'w') as f: |
|
|
json.dump(params, f) |
|
|
|
|
|
|
|
|
device = params['device'] |
|
|
run_tests = params['run_tests'] |
|
|
save_results = params['save_results'] |
|
|
dataset_path = params['dataset_path'] |
|
|
|
|
|
|
|
|
train_dataset_size = params['train_dataset_size'] |
|
|
validation_dataset_size = params['validation_dataset_size'] |
|
|
baseline_dataset_size = params['baseline_dataset_size'] |
|
|
|
|
|
|
|
|
num_nodes = params['num_nodes'] |
|
|
num_depots = params['num_depots'] |
|
|
embedding_size = params['embedding_size'] |
|
|
sample_size = params['sample_size'] |
|
|
gradient_clipping = params['gradient_clipping'] |
|
|
num_neighbors_encoder = params['num_neighbors_encoder'] |
|
|
num_neighbors_action = params['num_neighbors_action'] |
|
|
num_movers = params['num_movers'] |
|
|
learning_rate = params['learning_rate'] |
|
|
batch_size = params['batch_size'] |
|
|
test_batch_size = params['test_batch_size'] |
|
|
|
|
|
|
|
|
validation_dataset = VRP_Dataset(validation_dataset_size, num_nodes, num_depots, dataset_path, device) |
|
|
baseline_dataset = VRP_Dataset(train_dataset_size, num_nodes, num_depots, dataset_path, device) |
|
|
if params['overfit_test']: |
|
|
train_dataset = baseline_dataset = validation_dataset = VRP_Dataset(train_dataset_size, num_nodes, num_depots, dataset_path, device) |
|
|
|
|
|
|
|
|
google_scores = evaluate_google_model(validation_dataset) |
|
|
tot_google_scores = google_scores.sum().item() |
|
|
input_size = validation_dataset.model_input_length() |
|
|
|
|
|
|
|
|
model = Model(input_size=input_size, embedding_size=embedding_size, decoder_input_size=params["decoder_input_size"]) |
|
|
actor = Actor(model=model, num_movers=num_movers, num_neighbors_encoder=num_neighbors_encoder, |
|
|
num_neighbors_action=num_neighbors_action, device=device, normalize=False) |
|
|
actor.train_mode() |
|
|
|
|
|
baseline_model = Model(input_size=input_size, embedding_size=embedding_size, decoder_input_size=params["decoder_input_size"]) |
|
|
baseline_actor = Actor(model=baseline_model, num_movers=num_movers, num_neighbors_encoder=num_neighbors_encoder, |
|
|
num_neighbors_action=num_neighbors_action, device=device, normalize=False) |
|
|
baseline_actor.greedy_search() |
|
|
baseline_actor.load_state_dict(actor.state_dict()) |
|
|
|
|
|
nn_actor = Actor(model=None, num_movers=1, num_neighbors_action=1, device=device) |
|
|
nn_actor.nearest_neighbors() |
|
|
|
|
|
optimizer = optim.Adam(params=actor.parameters(), lr=learning_rate) |
|
|
|
|
|
|
|
|
train_results_file = os.path.join(safe_data_dir, "train_results.txt") |
|
|
test_results_file = os.path.join(safe_data_dir, "test_results.txt") |
|
|
model_path = os.path.join(safe_data_dir, "model_state_dict.pt") |
|
|
optimizer_path = os.path.join(safe_data_dir, "optimizer_state_dict.pt") |
|
|
|
|
|
train_batch_record = 100 |
|
|
validation_record = 100 |
|
|
baseline_record = None |
|
|
|
|
|
|
|
|
for epoch in range(params['num_epochs']): |
|
|
if not params['overfit_test']: |
|
|
train_dataset = VRP_Dataset(train_dataset_size, num_nodes, num_depots, dataset_path, device) |
|
|
|
|
|
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, collate_fn=train_dataset.collate) |
|
|
for i, batch in enumerate(train_dataloader): |
|
|
with torch.no_grad(): |
|
|
nn_output = nn_actor(batch) |
|
|
tot_nn_cost = nn_output['total_time'].sum().item() |
|
|
|
|
|
baseline_actor.greedy_search() |
|
|
baseline_cost = baseline_actor(batch)['total_time'] |
|
|
|
|
|
actor.train_mode() |
|
|
actor_output = actor(batch) |
|
|
actor_cost, log_probs = actor_output['total_time'], actor_output['log_probs'] |
|
|
|
|
|
loss = ((actor_cost - baseline_cost).detach() * log_probs).mean() |
|
|
|
|
|
optimizer.zero_grad() |
|
|
loss.backward() |
|
|
if gradient_clipping: |
|
|
for group in optimizer.param_groups: |
|
|
clip_grad_norm_(group['params'], 1, norm_type=2) |
|
|
optimizer.step() |
|
|
|
|
|
tot_actor_cost = actor_cost.sum().item() |
|
|
tot_baseline_cost = baseline_cost.sum().item() |
|
|
actor_nn_ratio = tot_actor_cost / tot_nn_cost |
|
|
actor_baseline_ratio = tot_actor_cost / tot_baseline_cost |
|
|
train_batch_record = min(train_batch_record, actor_nn_ratio) |
|
|
|
|
|
result = f"{epoch}, {i}, {actor_nn_ratio:.4f}, {actor_baseline_ratio:.4f}, {train_batch_record:.4f}" |
|
|
print(result) |
|
|
if save_results: |
|
|
with open(train_results_file, 'a') as f: |
|
|
f.write(result + '\n') |
|
|
del batch |
|
|
|
|
|
|
|
|
if epoch % 5 == 0: |
|
|
baseline_dataloader = DataLoader(baseline_dataset, batch_size=batch_size, collate_fn=baseline_dataset.collate) |
|
|
tot_cost = [] |
|
|
for batch in baseline_dataloader: |
|
|
with torch.no_grad(): |
|
|
actor.greedy_search() |
|
|
cost = actor(batch)['total_time'] |
|
|
tot_cost.append(cost) |
|
|
tot_cost = torch.cat(tot_cost, dim=0) |
|
|
if baseline_record is None or (tot_cost < baseline_record).float().mean().item() > 0.9: |
|
|
baseline_record = tot_cost |
|
|
baseline_actor.load_state_dict(actor.state_dict()) |
|
|
print('\nNew baseline record\n') |
|
|
|
|
|
|
|
|
if (epoch % 10 == 0) and run_tests: |
|
|
b = max(int(batch_size // sample_size**2), 1) |
|
|
validation_dataloader = DataLoader(validation_dataset, batch_size=b, collate_fn=validation_dataset.collate) |
|
|
|
|
|
tot_cost = 0 |
|
|
tot_nn_cost = 0 |
|
|
for batch in validation_dataloader: |
|
|
with torch.no_grad(): |
|
|
actor.beam_search(sample_size) |
|
|
actor_output = actor(batch) |
|
|
cost = actor_output['total_time'] |
|
|
|
|
|
nn_output = nn_actor(batch) |
|
|
nn_cost = nn_output['total_time'] |
|
|
|
|
|
tot_cost += cost.sum().item() |
|
|
tot_nn_cost += nn_cost.sum().item() |
|
|
|
|
|
ratio = tot_cost / tot_nn_cost |
|
|
validation_record = min(validation_record, ratio) |
|
|
actor_google_ratio = tot_cost / tot_google_scores |
|
|
print(f"\nTest results:\nActor/Google: {actor_google_ratio:.4f}, Actor/NN: {ratio:.4f}, Best NN Ratio: {validation_record:.4f}\n") |
|
|
|
|
|
if save_results: |
|
|
with open(test_results_file, 'a') as f: |
|
|
f.write(f"{epoch}, {actor_google_ratio:.4f}, {ratio:.4f}, {validation_record:.4f}\n") |
|
|
torch.save(actor.state_dict(), model_path) |
|
|
torch.save(optimizer.state_dict(), optimizer_path) |
|
|
|
|
|
print("End.") |
|
|
|
