def main(config):
# 选择使用模式
if config.mode == 'train':
# 获取DataLoader
train_loader = get_loader(config)
# 计数, 找到可以创建文件夹的名字
run = 0
while os.path.exists("%s/run-%d" % (config.save_folder, run)):
run += 1
# 创建保存模型文件夹
os.mkdir("%s/run-%d" % (config.save_folder, run))
os.mkdir("%s/run-%d/models" % (config.save_folder, run))
# 在config中保存文件夹名称
config.save_folder = "%s/run-%d" % (config.save_folder, run)
# 开始进入training状态
train = Solver(train_loader, None, config)
train.train()
elif config.mode == 'test':
# 选择数据集路径
config.test_root, config.test_list = get_test_info(config.sal_mode)
# 获取DataLoader
test_loader = get_loader(config, mode='test')
# 如果文件路径不存在
if not os.path.exists(config.test_fold):
os.mkdir(config.test_fold)
# 开始进入testing状态
test = Solver(None, test_loader, config)
test.test()
else:
# 非法输入
raise IOError("illegal input!!!")
def main(config):
cudnn.benchmark = True
if config.model_type not in ['U_Net','FTC_Net']:
print('ERROR!! model_type should be selected in U_Net/FTC_Net')
print('Your input for model_type was %s'%config.model_type)
return
if not os.path.exists(config.model_path):
os.mkdir(config.model_path)
if not os.path.exists(config.result_path):
os.mkdir(config.result_path)
config.result_path = os.path.join(config.result_path,config.model_type)
if not os.path.exists(config.result_path):
os.mkdir(config.result_path)
train_loader = get_loader(config.save_path,'train',config.batch_size,config.num_workers)
test_loader = get_loader(config.save_path, 'test', 1, config.num_workers)
print('train dataset len %d' % (len(train_loader) * config.batch_size))
print('test dataset len %d' % len(test_loader))
results_path = config.result_path
layer_results_path = os.path.join(results_path, 'layers')
dst_results_path = os.path.join(results_path, 'dst_layer')
att_result_path = os.path.join(results_path, 'attention_map')
confi_result_path = os.path.join(results_path, 'confidence_map')
if not os.path.exists(results_path):
os.mkdir(results_path)
if not os.path.exists(layer_results_path):
os.mkdir(layer_results_path)
os.mkdir(dst_results_path)
if not os.path.exists(att_result_path):
os.mkdir(att_result_path)
if not os.path.exists(confi_result_path):
os.mkdir(confi_result_path)
if config.flag == 'FTC_Net':
solver = Solver(config, train_loader, test_loader,)
if config.mode == 'train':
solver.train()
elif config.mode == 'test':
# os.mkdir(contour_results_path)
# os.mkdir(attmap_path)
# solver.test(layer_results_path,dst_results_path,att_result_path,config.is_save_attsmat,'./meter')
solver.test(layer_results_path)
elif config.mode == 'confi':
solver.cal_confi(confi_result_path)
else:
solver.get_atts(layer_results_path,dst_results_path,att_result_path,config.is_save_attsmat)
def export(net,
setup,
system_env,
system,
exported_path="exported/output/net.pt",
batch_size=1):
print(colored("Trying to export to {}".format(exported_path), "red"))
env_vox = torch.from_numpy(
np.load('{}/data/{}_env_vox.npy'.format(system_env, system))).float()
train_loader, test_loader = get_loader(system_env,
system,
batch_size=1,
setup=setup)
env_input = env_vox[train_loader.dataset[0:batch_size][0]
[:, 0].long()].cuda()
state_input = train_loader.dataset[0:batch_size][0][:, 1:].cuda()
eval_net = copy.deepcopy(net)
eval_net.eval()
output = eval_net(state_input, env_input)
print(colored("eval_output:", "blue"))
print(output)
eval_serilized_module = torch.jit.script(eval_net)
print(colored("eval_serilized_output:", "blue"))
print(eval_serilized_module(state_input, env_input))
print(colored("Tracing and Saving...", "blue"))
traced_script_module = torch.jit.trace(net, (state_input, env_input))
serilized_module = torch.jit.script(net)
serilized_output = serilized_module(state_input, env_input)
traced_script_module.save(exported_path)
print(colored("Exported to {}".format(exported_path), "red"))
def main(ae_output_size, state_size, lr, epochs, batch, system_env, system,
setup, loss_type):
mpnet = VAEMPNet(ae_input_size=32,
ae_output_size=1024,
in_channels=1,
state_size=4,
code_size=32)
data_loaders = get_loader(system_env,
system,
batch_size=batch,
setup=setup)
train_network(network=mpnet,
data_loaders=data_loaders,
network_name="mpnet_vae",
lr=lr,
epochs=epochs,
batch=batch,
system_env=system_env,
system=system,
setup=setup,
using_step_lr=True,
step_size=50,
gamma=0.9,
loss_type=loss_type,
weight_save_epochs=50)
def train(lr, epochs, batch, system, model, setup):
mpnet = MPNet(ae_input_size=32,
ae_output_size=128,
in_channels=1,
state_size=4,
control_size=2)
env_vox = torch.from_numpy(
np.load('{}/{}_env_vox.npy'.format(system, model))).float()
if torch.cuda.is_available():
mpnet = mpnet.cuda()
env_vox = env_vox.cuda()
#optimizer = torch.optim.Adagrad(mpnet.parameters(), lr=lr)
#optimizer = torch.optim.SGD(mpnet.parameters(), lr=lr, momentum=0.9)
optimizer = torch.optim.Adam(mpnet.parameters(), lr=lr)
#scheduler = StepLR(optimizer, step_size=10, gamma=0.9)
train_loader, test_loader = get_loader(system,
model,
batch_size=batch,
setup=setup)
for i in tqdm(range(epochs + 1)):
print("epoch {}".format(i))
train_loss = []
mpnet.train()
for data, label in tqdm(train_loader):
## prepare data
if torch.cuda.is_available():
data = data.cuda()
label = label.cuda()
inputs = data[:, 1:]
envs = env_vox[(data[:, 0]).long()]
## execute
optimizer.zero_grad()
output = mpnet(inputs, envs)
loss = F.mse_loss(output, label)
loss.backward()
optimizer.step()
train_loss.append(loss.item())
# break
#scheduler.step(i)
print("train loss:{}".format(np.mean(train_loss)))
for data, label in test_loader:
mpnet.eval()
eval_loss = []
with torch.no_grad():
if torch.cuda.is_available():
data = data.cuda()
label = label.cuda()
inputs = data[:, 1:]
envs = env_vox[(data[:, 0]).long()]
output = mpnet(inputs, envs)
loss = F.mse_loss(output, label)
eval_loss.append(loss.item())
print("eval loss:{}".format(np.mean(eval_loss)))
if i % 5 == 0:
Path("output/{}/{}".format(model, setup)).mkdir(parents=True,
exist_ok=True)
torch.save(mpnet.state_dict(),
"output/{}/{}/ep{}".format(model, setup, i))
def main(config):
if config.mode == 'train':
train_loader = get_loader(config)
run = 0
while os.path.exists("%s/run-%d" % (config.save_folder, run)):
run += 1
os.mkdir("%s/run-%d" % (config.save_folder, run))
os.mkdir("%s/run-%d/models" % (config.save_folder, run))
config.save_folder = "%s/run-%d" % (config.save_folder, run)
train = Solver(train_loader, None, config)
train.train()
elif config.mode == 'test':
config.test_root, config.test_list = get_test_info(config.sal_mode)
test_loader = get_loader(config, mode='test')
if not os.path.exists(config.test_fold): os.mkdir(config.test_fold)
test = Solver(None, test_loader, config)
test.test()
else:
raise IOError("illegal input!!!")
def main(lr, num_epoch, batch_size,
loss_type, network_tag):
net = CNNWrapper(loss_type=loss_type, save_path="./output", device="cuda", tag=network_tag)
data_loaders = get_loader(batch_size=batch_size)
tvs = TrainValScheduler(network=net, data_loaders=data_loaders,
lr=1e-3, num_epoch=1, batch_size=200,
using_step_lr=False, step_size=50, gamma=0.9,
weight_save_epochs=50,
logging_path="log/", optimizer_type="Adadelta")
tvs.train_val()
def main(ae_output_size, state_size, lr, epochs, batch, system_env, system,
setup, loss_type, lr_step_size, aug, network_name):
# mpnet_module = importlib.import_module('.mpnet_{}'.format(system), package=".networks")
try:
if system == 'cartpole_obs':
if network_name == 'mpnet':
from networks.mpnet_cartpole_obs import MPNet
elif network_name == 'mpnet_branch':
from networks.mpnet_cartpole_obs_branch import MPNet
state_dim = 4
in_channels = 1
elif system == 'acrobot_obs':
from networks.mpnet_acrobot_obs import MPNet
state_dim = 4
in_channels = 1
elif system == 'quadrotor_obs':
from networks.mpnet_quadrotor_obs import MPNet
state_dim = 13
in_channels = 32
elif system == 'car_obs':
from networks.mpnet_car_obs import MPNet
state_dim = 3
in_channels = 1
except ModuleNotFoundError:
print("Unrecognized model name")
raise
mpnet = MPNet(ae_input_size=32,
ae_output_size=ae_output_size,
in_channels=in_channels,
state_size=state_dim)
data_loaders = get_loader(system_env,
system,
batch_size=batch,
setup=setup)
train_network(network=mpnet,
data_loaders=data_loaders,
network_name=network_name,
lr=lr,
epochs=epochs,
batch=batch,
system_env=system_env,
system=system,
setup=setup,
using_step_lr=True,
step_size=lr_step_size,
gamma=0.9,
loss_type=loss_type,
weight_save_epochs=25,
aug=aug)
def main(config):
test_loader = get_loader(test_mode=config.test_mode, sal_mode=config.sal_mode)
if not os.path.exists(config.test_fold): os.mkdir(config.test_fold)
test = Solver(test_loader, config)
test.test(test_mode=config.test_mode)
def train(lr, epochs, batch, system, model, setup):
mpnet = VAEMPNet(ae_input_size=32,
ae_output_size=1024,
in_channels=1,
state_size=4,
code_size=32)
env_vox = torch.from_numpy(
np.load('{}/{}_env_vox.npy'.format(system, model))).float()
if torch.cuda.is_available():
mpnet = mpnet.cuda()
env_vox = env_vox.cuda()
optimizer = torch.optim.Adam(mpnet.parameters(), lr=lr)
scheduler = StepLR(optimizer, step_size=50, gamma=0.9)
train_loader, test_loader = get_loader(system,
model,
batch_size=batch,
setup=setup)
with tqdm(range(epochs + 1), total=epochs + 1) as pbar:
for i in range(epochs + 1):
#print("epoch {}".format(i))
train_loss = []
mpnet.train()
for data, label in tqdm(train_loader):
## prepare data
if torch.cuda.is_available():
data = data.cuda()
label = label.cuda()
inputs = data[:, 1:]
envs = env_vox[(data[:, 0]).long()]
## execute
optimizer.zero_grad()
output = mpnet(inputs, envs)
# loss = F.mse_loss(output, label)
loss = F.l1_loss(output, label)
loss.backward()
# print(loss.item())
# train_loss.append(loss.item())
optimizer.step()
# break
scheduler.step(i)
# print("train loss:{}".format(np.mean(train_loss)))
mpnet.eval()
for data, label in test_loader:
eval_loss = []
with torch.no_grad():
if torch.cuda.is_available():
data = data.cuda()
label = label.cuda()
inputs = data[:, 1:]
envs = env_vox[(data[:, 0]).long()]
output = mpnet(inputs, envs)
# loss = F.mse_loss(output, label)
loss = F.l1_loss(output, label)
eval_loss.append(loss.item())
#print("eval loss:{}".format(np.mean(eval_loss)))
pbar.set_postfix(
{'eval_loss': '{0:1.5f}'.format(np.mean(eval_loss))})
if i % 50 == 0:
Path("output/{}/{}".format(model, setup)).mkdir(parents=True,
exist_ok=True)
torch.save(mpnet.state_dict(),
"output/{}/{}/ep{}".format(model, setup, i))
pbar.update(1)
parser.add_argument('--train_root', type=str, default='')
parser.add_argument('--train_list', type=str, default='')
# Testing settings
parser.add_argument('--model', type=str, default='./dataset/pretrained/final.pth') # Snapshot
parser.add_argument('--test_fold', type=str, default='./results') # Test results saving folder
parser.add_argument('--sal_mode', type=str, default='g') # Test image dataset
# Misc
parser.add_argument('--mode', type=str, default='test', choices=['train', 'test'])
config = parser.parse_args()
if not os.path.exists(config.save_folder):
os.mkdir(config.save_folder)
# Get test set info
config.test_root = './data/game/'
config.test_list = './data/game/test.lst'
config.cuda = False
test_loader = get_loader(config, mode='test')
if not os.path.exists(config.test_fold): os.mkdir(config.test_fold)
test = Solver(None, test_loader, config)
#test.test()
#####################################################
# (2) load pretrained weights
#####################################################
#------------------------------------------
import io
import numpy as np
from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx
import torch.nn as nn
