Homec4science
Diffusion QPFit (master) Browse

QPFit/run_fit.py
3bfb25b5e70amaster

Last Change
Raw File

run_fit.py
View Options

import sys
print(sys.version)
import numpy as np
import timeit
import torch
import scipy.io as sio
from qpfit import QPNet, GradNet
from qpfit import QPDataset
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import os
torch.set_num_threads(os.cpu_count())
print("Number of threads = " + str(torch.get_num_threads()))
param = dict()
param['ModelType'] = 'GradNet'
# param['ModelType'] = 'QPNet'
print(sys.argv)
param['nVar'] = int(sys.argv[1])
param['numItr'] = int(sys.argv[2])
print("==================================\n")
print("Running with nVar = %i and numItr = %i\n" % (param['nVar'], param['numItr']))
print("==================================\n")
param['maxIter'] = 10
param['batch_size'] = 50
param['num_epochs'] = 100
param['log_interval'] = 10000
param['learning_rate'] = 1e-3
datafile = 'double_integrator'
# datafile = 'fourstate'
# datafile = 'double_track'
# datafile = 'dcmotor_2000000'
# datafile = 'bandp_200000'
param['umax'] = 2
# param['umax'] = 220
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = 'cuda:0'
# device = 'cpu'
savefile = datafile + '_' + param['ModelType'] + '_' + 'var' + str(param['nVar']) + '_itr_' + str(param['numItr'])
qp_loader = QPDataset(datafile)
dataloader = DataLoader(qp_loader, batch_size=param['batch_size'],
shuffle=True, num_workers=torch.get_num_threads(), pin_memory=True)
# Initialize the model.
nParam = qp_loader.x.shape[1]
nInputs = qp_loader.y.shape[1]
if param['ModelType'] == 'QPNet':
model = QPNet(param['nVar'], in_features=nParam, out_features=nInputs,
maxIter=param['maxIter'], umax=param['umax'])
else:
model = GradNet(param['nVar'], in_features=nParam, out_features=nInputs,
numItr=param['numItr'], umax=param['umax'], device=device)
model = model.to(device)
loss_fn = torch.nn.MSELoss()
print(model.type)
# Confirm we can save the model
# model.save(datafile)
# TRAIN!
num_epochs = param['num_epochs']
log_interval = param['log_interval']
# Initialize the optimizer.
learning_rate = param['learning_rate']
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# optimizer = torch.optim.Adamax(model.parameters())
# optimizer = torch.optim.LBFGS(model.parameters())
print('Training initialized')
model.train()
for epoch in range(num_epochs):
for i_batch, (x, y) in enumerate(dataloader):
x, y = x.to(device), y.to(device)
def closure():
if torch.is_grad_enabled():
optimizer.zero_grad()
y_pred = model(x)
loss = loss_fn(y_pred, y)
if loss.requires_grad:
loss.backward()
return loss
optimizer.step(closure)
# x, y = x.to(device), y.to(device)
# optimizer.zero_grad()
# y_pred = model(x)
# loss = loss_fn(y_pred, y)
# loss.backward()
# optimizer.step()
if i_batch % log_interval == 0:
with torch.no_grad():
loss = closure()
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, i_batch * len(x), len(dataloader.dataset),
100. * i_batch / len(dataloader), loss.item()))
filename = savefile + '_epoch_' + str(epoch)
print("Saving file " + filename)
model.save(filename)
print('Training complete - loss = %f' % loss)
c4science · Help