def main(scripts, args):
scripts = " ".join(sys.argv[0:])
args = parser.parse_args()
args.scripts = scripts
torch.manual_seed(args.seed)
if args.gpu_ids == -1:
args.gpu_ids = [-1]
else:
torch.cuda.manual_seed(args.seed)
mp.set_start_method('spawn')
if (args.deploy):
raw, gt_lbl, raw_valid, gt_lbl_valid, raw_test, gt_lbl_test, raw_test_upsize, gt_lbl_test_upsize = setup_data(
args)
else:
raw, gt_lbl, raw_valid, gt_lbl_valid, raw_test, gt_lbl_test = setup_data(
args)
env_conf = setup_env_conf(args)
shared_model = get_model(args,
args.model,
env_conf["observation_shape"],
args.features,
atrous_rates=args.atr_rate,
num_actions=2,
split=args.data_channel,
multi=args.multi)
manager = mp.Manager()
shared_dict = manager.dict()
if args.wctrl == "s2m":
shared_dict["spl_w"] = args.spl_w
shared_dict["mer_w"] = args.mer_w
if args.load:
saved_state = torch.load(args.load,
map_location=lambda storage, loc: storage)
shared_model.load_state_dict(saved_state)
shared_model.share_memory()
if args.shared_optimizer:
if args.optimizer == 'RMSprop':
optimizer = SharedRMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = SharedAdam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
optimizer.share_memory()
else:
optimizer = None
processes = []
if not args.no_test:
if raw_test is not None:
if (args.deploy):
p = mp.Process(target=test_func,
args=(args, shared_model, env_conf,
[raw_valid, gt_lbl_valid],
(raw_test, gt_lbl_test, raw_test_upsize,
gt_lbl_test_upsize, shared_dict)))
else:
p = mp.Process(target=test_func,
args=(args, shared_model, env_conf,
[raw_valid, gt_lbl_valid
], (raw_test, gt_lbl_test), shared_dict))
else:
p = mp.Process(target=test_func,
args=(args, shared_model, env_conf,
[raw_valid, gt_lbl_valid], None, shared_dict))
p.start()
processes.append(p)
time.sleep(0.1)
for rank in range(0, args.workers):
p = mp.Process(target=train_func,
args=(rank, args, shared_model, optimizer, env_conf,
[raw, gt_lbl], shared_dict))
p.start()
processes.append(p)
time.sleep(0.1)
for p in processes:
time.sleep(0.1)
p.join()
if args.load:
saved_state = torch.load('{0}{1}.dat'.format(args.load_model_dir,
args.env),
map_location=lambda storage, loc: storage)
shared_model.load_state_dict(saved_state)
shared_model.share_memory()
if args.shared_optimizer:
if args.optimizer == 'RMSprop':
optimizer = SharedRMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = SharedAdam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
optimizer.share_memory()
else:
optimizer = None
processes = []
if "EM_env" in args.env:
p = mp.Process(target=test,
args=(args, shared_model, env_conf,
[raw, lbl, prob, gt_lbl], True))
else:
p = mp.Process(target=test, args=(args, shared_model, env_conf))
p.start()
processes.append(p)
time.sleep(1)
if "EM_env" in args.env:
def start():
args = parser.parse_args()
args.shared_optimizer = True
if args.gpu_ids == -1:
torch.manual_seed(args.seed)
args.gpu_ids = [-1]
device_share = torch.device('cpu')
mp.set_start_method('spawn')
else:
torch.cuda.manual_seed(args.seed)
mp.set_start_method('spawn', force=True)
if len(args.gpu_ids) > 1:
device_share = torch.device('cpu')
else:
device_share = torch.device('cuda:' + str(args.gpu_ids[-1]))
env = create_env(args.env, args)
shared_model = build_model(env.observation_space, env.action_space, args,
device_share).to(device_share)
shared_model.share_memory()
env.close()
del env
if args.load_coordinator_dir is not None:
saved_state = torch.load(args.load_coordinator_dir,
map_location=lambda storage, loc: storage)
if args.load_coordinator_dir[-3:] == 'pth':
shared_model.load_state_dict(saved_state['model'], strict=False)
else:
shared_model.load_state_dict(saved_state)
params = shared_model.parameters()
if args.shared_optimizer:
print('share memory')
if args.optimizer == 'RMSprop':
optimizer = SharedRMSprop(params, lr=args.lr)
if args.optimizer == 'Adam':
optimizer = SharedAdam(params, lr=args.lr, amsgrad=args.amsgrad)
optimizer.share_memory()
else:
optimizer = None
current_time = datetime.now().strftime('%b%d_%H-%M')
args.log_dir = os.path.join(args.log_dir, args.env, current_time)
processes = []
manager = mp.Manager()
train_modes = manager.list()
n_iters = manager.list()
p = mp.Process(target=test,
args=(args, shared_model, optimizer, train_modes, n_iters))
p.start()
processes.append(p)
time.sleep(args.sleep_time)
for rank in range(0, args.workers):
p = mp.Process(target=train,
args=(rank, args, shared_model, optimizer, train_modes,
n_iters))
p.start()
processes.append(p)
time.sleep(args.sleep_time)
for p in processes:
time.sleep(args.sleep_time)
p.join()
def run(self):
torch.manual_seed(args.seed)
if args.gpu_ids == -1:
args.gpu_ids = [-1]
else:
torch.cuda.manual_seed(args.seed)
mp.set_start_method('spawn')
# env = make_env(env_type=args.env_type, env_name=args.env_name, args=args)
shared_model = UNREAL(in_channels=3, action_size=6, enable_pixel_control=True)
if args.load:
saved_state = torch.load(
'{0}{1}.dat'.format(args.load_model_dir, args.env),
map_location=lambda storage, loc: storage)
shared_model.load_state_dict(saved_state)
shared_model.share_memory()
lr = log_uniform(1e-4, 5e-3, 0.5)
if args.shared_optimizer:
if args.optimizer == 'RMSprop':
optimizer = SharedRMSprop(shared_model.parameters(), lr=lr, eps=0.1)
if args.optimizer == 'Adam':
optimizer = SharedAdam(
shared_model.parameters(), lr=lr, amsgrad=args.amsgrad)
optimizer.share_memory()
else:
optimizer = None
# p = mp.Process(target=train, args=(args, shared_model, env_conf))
# p.start()
# processes.append(p)
# time.sleep(0.1)
self.stop_requested = False
self.terminate_reqested = False
for rank in range(0, args.workers):
trainer = Trainer(rank, args, shared_model=shared_model, optimizer=optimizer, lr=lr)
self.trainers.append(trainer)
# time.sleep(0.1)
# 设置运行起始的时间
# set start time
self.start_time = time.time() - 0 # wall_t
processes = []
for rank in range(0, args.workers):
if rank == 0:
p = mp.Process(target=self.train_function, args=(rank, True, True))
else:
p = mp.Process(target=self.train_function, args=(rank, True))
p.start()
processes.append(p)
# 注册终止信号
signal.signal(signal.SIGINT, self.signal_handler)
print('Press Ctrl+C to stop')
for rank in range(0, args.workers):
time.sleep(0.01)
processes[rank].join()
if __name__ == "__main__":
params = Params()
mp.set_start_method('spawn')
count = mp.Value('i', 0) # update count
lock = mp.Lock()
# shared_model = A3C()
shared_model = A3C_LSTM()
shared_model = shared_model.share_memory()
# shared_optimizer = SharedAdam(shared_model.parameters(), lr=params.lr, amsgrad=params.amsgrad,
# weight_decay=params.weight_decay)
shared_optimizer = SharedRMSprop(shared_model.parameters(), lr=params.lr)
shared_optimizer.share_memory()
# run_loop(0, params, shared_model, shared_optimizer, count, lock) # for debugging
# test(0, params, shared_model, count, lock)
processes = []
# have to add test module
p = mp.Process(target=test, args=(
0,
params,
shared_model,
count,
lock,
))
p.start()