def __init__(
self,
_agent: AgentAbstract,
_env: EnvAbstract,
_epoch_max: int,
_epoch_train: int,
_train_update_target: int,
_train_save: int,
_process_core: int = None,
_save_path: str = './save',
_use_cmd: bool = True,
):
self.agent: AgentAbstract = _agent
self.agent.training()
self.env: EnvAbstract = _env
# multiprocessing for sampling
self.mp = mp.get_context('spawn')
self.process_core = _process_core
self.pool = self.mp.Pool(self.process_core)
# training control
self.epoch = 0
self.train_times = 0
self.epoch_max = _epoch_max
self.epoch_train = _epoch_train
self.train_update_target = _train_update_target
self.train_save = _train_save
self.total_reward_buf = []
self.save_path = _save_path
self.use_cmd = _use_cmd
if self.use_cmd:
self.shell = TrainShell(self)
def mp_search(self, graph, other_info, model_id, train_data, test_data):
ctx = mp.get_context()
q = ctx.Queue()
p = ctx.Process(target=train, args=(q, graph, train_data, test_data, self.trainer_args,
self.metric, self.loss, self.verbose, self.path))
try:
p.start()
search_results = self._search_common(q)
metric_value, loss, graph = q.get(block=True)
if time.time() >= self._timeout:
raise TimeoutError
if self.verbose and search_results:
for (generated_graph, generated_other_info, new_model_id) in search_results:
verbose_print(generated_other_info, generated_graph, new_model_id)
if metric_value is not None:
self.add_model(metric_value, loss, graph, model_id)
self.update(other_info, model_id, graph, metric_value)
except (TimeoutError, queue.Empty) as e:
raise TimeoutError from e
finally:
# terminate and join the subprocess to prevent any resource leak
p.terminate()
p.join()
def test_cuda_small_tensors(self):
# Check multiple small tensors which will likely use the same
# underlying cached allocation
ctx = mp.get_context('spawn')
tensors = []
for i in range(5):
device = i % 2
tensors += [torch.arange(i * 5, (i + 1) * 5).cuda(device)]
inq = ctx.Queue()
outq = ctx.Queue()
inq.put(tensors)
p = ctx.Process(target=sum_tensors, args=(inq, outq))
p.start()
results = []
for i in range(5):
results.append(outq.get())
p.join()
for i, tensor in enumerate(tensors):
v, device, tensor_size, storage_size = results[i]
self.assertEqual(v, torch.arange(i * 5, (i + 1) * 5).sum())
self.assertEqual(device, i % 2)
self.assertEqual(tensor_size, 5)
self.assertEqual(storage_size, 5)
def __init__(self, device_ids=None, multiprocessing_method='spawn'):
super().__init__()
self.device_ids = tuple(device_ids)
self.num_replicas = len(device_ids)
self.rank = None
self._mp = multiprocessing.get_context(multiprocessing_method)
self._start_error_handler()
self._start_multiprocessing()
def __init__(self, loader, prepro,
sort_key, batchify,
single_run=True, queue_size=8, fork=True):
self._loader = loader
self._prepro = prepro
self._sort_key = sort_key
self._batchify = batchify
self._single_run = single_run
if fork:
ctx = mp.get_context('forkserver')
self._queue = ctx.Queue(queue_size)
else:
# for easier debugging
self._queue = None
self._process = None
def test_event(self):
ctx = mp.get_context('spawn')
queue = ctx.Queue()
ready = ctx.Event()
done = ctx.Event()
p = ctx.Process(target=cuda_multiply_two, args=(queue, ready, done))
p.start()
ready.wait()
with torch.cuda.stream(torch.cuda.Stream()):
tensor = torch.cuda.FloatTensor([1, 1, 1, 1])
# Use a sleep kernel to test events. Without the event, the
# multiply happens before the add.
event = torch.cuda.Event(interprocess=True)
torch.cuda._sleep(20000000) # about 30 ms
tensor.add_(1)
event.record()
queue.put((event, tensor))
done.wait() # must wait until subprocess records event
event.synchronize()
self.assertEqual(list(tensor), [4, 4, 4, 4])
p.join()
def __call__(self, batch_size: int):
def get_batches(hyper_batch):
indexes = list(range(0, len(hyper_batch), batch_size))
if not self._single_run:
# random shuffle for training batches
random.shuffle(hyper_batch)
random.shuffle(indexes)
hyper_batch.sort(key=self._sort_key)
for i in indexes:
batch = self._batchify(hyper_batch[i:i+batch_size])
yield batch
if self._queue is not None:
ctx = mp.get_context('forkserver')
self._process = ctx.Process(
target=_batch2q,
args=(self._loader, self._prepro,
self._queue, self._single_run)
)
self._process.start()
while True:
d = self._queue.get()
if d is None:
break
if isinstance(d, int):
print('\nepoch {} done'.format(d))
continue
yield from get_batches(d)
self._process.join()
else:
i = 0
while True:
for batch in self._loader:
yield from get_batches(self._prepro(batch))
if self._single_run:
break
i += 1
print('\nepoch {} done'.format(i))
def setup(self, model):
os.environ["MASTER_PORT"] = str(self.cluster_environment.master_port())
# pass in a state q
smp = mp.get_context("spawn")
self.mp_queue = smp.SimpleQueue()
def __init__(
self,
make_sampler_fn: Callable[..., TaskSampler],
sampler_fn_args: Sequence[Dict[str, Any]] = None,
auto_resample_when_done: bool = True,
multiprocessing_start_method: Optional[str] = "forkserver",
mp_ctx: Optional[BaseContext] = None,
metrics_out_queue: mp.Queue = None,
should_log: bool = True,
max_processes: Optional[int] = None,
) -> None:
self._is_waiting = False
self._is_closed = True
self.should_log = should_log
self.max_processes = max_processes
assert (
sampler_fn_args is not None and len(sampler_fn_args) > 0
), "number of processes to be created should be greater than 0"
self._num_task_samplers = len(sampler_fn_args)
self._num_processes = (
self._num_task_samplers
if max_processes is None
else min(max_processes, self._num_task_samplers)
)
self._auto_resample_when_done = auto_resample_when_done
assert (multiprocessing_start_method is None) != (
mp_ctx is None
), "Exactly one of `multiprocessing_start_method`, and `mp_ctx` must be not None."
if multiprocessing_start_method is not None:
assert multiprocessing_start_method in self._valid_start_methods, (
"multiprocessing_start_method must be one of {}. Got '{}'"
).format(self._valid_start_methods, multiprocessing_start_method)
self._mp_ctx = mp.get_context(multiprocessing_start_method)
else:
self._mp_ctx = typing.cast(BaseContext, mp_ctx)
self.metrics_out_queue = metrics_out_queue or self._mp_ctx.Queue()
self.npaused_per_process = [0] * self._num_processes
self.sampler_index_to_process_ind_and_subprocess_ind: Optional[
List[List[int]]
] = None
self._reset_sampler_index_to_process_ind_and_subprocess_ind()
self._workers: Optional[List] = None
for args in sampler_fn_args:
args["mp_ctx"] = self._mp_ctx
(
self._connection_read_fns,
self._connection_write_fns,
) = self._spawn_workers( # noqa
make_sampler_fn=make_sampler_fn,
sampler_fn_args_list=[
args_list for args_list in self._partition_to_processes(sampler_fn_args)
],
)
self._is_closed = False
for write_fn in self._connection_write_fns:
write_fn((OBSERVATION_SPACE_COMMAND, None))
observation_spaces = [
space for read_fn in self._connection_read_fns for space in read_fn()
]
if any(os is None for os in observation_spaces):
raise NotImplementedError(
"It appears that the `all_observation_spaces_equal`"
" is not True for some task sampler created by"
" VectorSampledTasks. This is not currently supported."
)
if any(observation_spaces[0] != os for os in observation_spaces):
raise NotImplementedError(
"It appears that the observation spaces of the samplers"
" created in VectorSampledTasks are not equal."
" This is not currently supported."
)
self.observation_space = observation_spaces[0]
for write_fn in self._connection_write_fns:
write_fn((ACTION_SPACE_COMMAND, None))
self.action_spaces = [
space for read_fn in self._connection_read_fns for space in read_fn()
]
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('-e', '--epochs', default='last', type=str)
parser.add_argument('-d', '--devices', default='1', type=str)
parser.add_argument('-v', '--verbose', default=False, action='store_true')
parser.add_argument('--show_image',
'-s',
default=False,
action='store_true')
parser.add_argument('--save_path', '-p', default=None)
args = parser.parse_args()
all_dev = parse_devices(args.devices)
mp_ctx = mp.get_context('spawn')
network = CPNet(config.num_classes, criterion=None)
data_setting = {
'img_root': config.img_root_folder,
'gt_root': config.gt_root_folder,
'train_source': config.train_source,
'eval_source': config.eval_source
}
dataset = ADE(data_setting, 'val', None)
with torch.no_grad():
segmentor = SegEvaluator(dataset, config.num_classes,
config.image_mean, config.image_std, network,
config.eval_scale_array, config.eval_flip,
all_dev, args.verbose, args.save_path,
args.show_image)
import torch
from torch import nn
from torch import optim
from torch import autograd
from torch import multiprocessing
from torch.autograd import Variable
from torch.nn import functional as F
from torch.utils.data import TensorDataset
import numpy as np
import os
multiprocessing = multiprocessing.get_context('spawn')
class Concept(nn.Module):
def __init__(self):
super(Concept, self).__init__()
self.net = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, padding=0),
nn.BatchNorm2d(64, momentum=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(64, 64, kernel_size=3, padding=0),
nn.BatchNorm2d(64, momentum=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64, momentum=1),
nn.ReLU(inplace=True),
nn.Conv2d(64, 64, kernel_size=3, padding=1),
def train(self):
if self.multiprocessing and self.view_id < 0:
import torch.multiprocessing as mp
mp = mp.get_context('spawn')
process = []
start = time.time()
for v in range(self.n_view):
process.append(mp.Process(target=self.train_view, args=(v, )))
process[v].daemon = True
process[v].start()
for p in process:
p.join()
elif self.view_id >= 0:
start = time.time()
self.train_view(self.view_id)
else:
start = time.time()
for v in range(self.n_view):
self.train_view(v)
end = time.time()
runing_time = end - start
if self.runing_time:
print('runing_time: ' + str(runing_time))
return runing_time
if not self.just_valid:
# valid = [self.resutls[v][0] for v in range(self.n_view)]
# test = [self.resutls[v][1] for v in range(self.n_view)]
valid_fea, valid_lab, test_fea, test_lab = [], [], [], []
for v in range(self.n_view):
tmp = sio.loadmat('features/' + self.datasets + '_' + str(v) +
'.mat')
valid_fea.append(tmp['valid_fea'])
valid_lab.append(tmp['valid_lab'].reshape([
-1,
]) if min(tmp['valid_lab'].shape) == 1 else tmp['valid_lab'])
test_fea.append(tmp['test_fea'])
test_lab.append(tmp['test_lab'].reshape([
-1,
]) if min(tmp['test_lab'].shape) == 1 else tmp['test_lab'])
valid_results = utils.multi_test(valid_fea, valid_lab, self.MAP)
test_results = utils.multi_test(test_fea, test_lab, self.MAP)
print("valid results: " + self.view_result(valid_results) +
",\t test resutls:" + self.view_result(test_results))
sio.savemat(
'features/' + self.datasets + '_SDML_test_feature_results.mat',
{
'test': test_fea,
'test_labels': test_lab
})
return valid_results, test_results
else:
return np.concatenate(
[np.array(loss).reshape([1, -1]) for loss in self.val_d_loss],
axis=0), np.concatenate([
np.array(loss).reshape([1, -1]) for loss in self.tr_d_loss
],
axis=0), np.concatenate([
np.array(loss).reshape([1, -1])
for loss in self.tr_ae_loss
],
axis=0)
def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
r"""Applies each `module` in :attr:`modules` in parallel on arguments
contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)
on each of :attr:`devices`.
Args:
modules (Module): modules to be parallelized
inputs (tensor): inputs to the modules
devices (list of int or torch.device): CUDA devices
:attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and
:attr:`devices` (if given) should all have same length. Moreover, each
element of :attr:`inputs` can either be a single object as the only argument
to a module, or a collection of positional arguments.
"""
assert len(modules) == len(inputs)
if kwargs_tup is not None:
assert len(modules) == len(kwargs_tup)
else:
kwargs_tup = ({}, ) * len(modules)
if devices is not None:
assert len(modules) == len(devices)
else:
devices = [None] * len(modules)
devices = list(map(lambda x: _get_device_index(x, True), devices))
context = mp.get_context('spawn')
# lock = threading.Lock()
# results = {}
# results = []
results_queue = context.Queue(len(devices))
grad_enabled = torch.is_grad_enabled()
def _worker(i, module, input, kwargs, device=None):
torch.set_grad_enabled(grad_enabled)
if device is None:
device = get_a_var(input).get_device()
try:
with torch.cuda.device(device):
# this also avoids accidental slicing of `input` if it is a Tensor
if not isinstance(input, (list, tuple)):
input = (input, )
output = module(*input, **kwargs)
results_queue.put(output)
# with lock:
# results[i] = output
except Exception as e:
results_queue.put(e)
# with lock:
# results[i] = e
if len(modules) > 1:
processes = [
context.Process(target=_worker,
args=(i, module, input, kwargs, device))
for i, (
module, input, kwargs,
device) in enumerate(zip(modules, inputs, kwargs_tup, devices))
]
for process in processes:
process.start()
for process in processes:
process.join()
else:
_worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])
outputs = []
for i in range(len(inputs)):
output = results_queue.get()
if isinstance(output, Exception):
raise output
outputs.append(output)
return outputs
def test_cuda_parameter_sharing(self):
param = Parameter(torch.arange(1., 26, device='cuda').view(5, 5))
self._test_autograd_sharing(param, mp.get_context('spawn'), is_parameter=True)
def test_cuda(self):
torch.cuda.FloatTensor([1]) # initialize CUDA outside of leak checker
self._test_sharing(mp.get_context('spawn'), torch.cuda.FloatTensor)
def train(flags): # pylint: disable=too-many-branches, too-many-statements
if flags.xpid is None:
flags.xpid = 'torchbeast-%s' % time.strftime('%Y%m%d-%H%M%S')
plogger = file_writer.FileWriter(
xpid=flags.xpid,
xp_args=flags.__dict__,
rootdir=flags.savedir,
symlink_latest=False,
)
checkpointpath = os.path.expandvars(
os.path.expanduser('%s/%s/%s' %
(flags.savedir, flags.xpid, 'model.tar')))
T = flags.unroll_length
B = flags.batch_size
flags.device = None
if not flags.disable_cuda and torch.cuda.is_available():
logging.info(f'Using CUDA {flags.cuda_id}')
flags.device = torch.device(f'cuda:{flags.cuda_id}')
else:
logging.info('Not using CUDA.')
flags.device = torch.device('cpu')
env = Net.create_env(flags)
model = Net.make(flags, env)
buffers = create_buffers(env.observation_space, len(env.action_space),
flags)
model.share_memory()
actor_processes = []
ctx = mp.get_context('fork')
free_queue = ctx.SimpleQueue()
full_queue = ctx.SimpleQueue()
for i in range(flags.num_actors):
actor = ctx.Process(target=act,
args=(i, free_queue, full_queue, model, buffers,
flags))
actor.start()
actor_processes.append(actor)
learner_model = Net.make(flags, env).to(device=flags.device)
optimizer = torch.optim.RMSprop(learner_model.parameters(),
lr=flags.learning_rate,
momentum=flags.momentum,
eps=flags.epsilon,
alpha=flags.alpha)
def lr_lambda(epoch):
return 1 - min(epoch * T * B, flags.total_frames) / flags.total_frames
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
if flags.resume:
save = torch.load(flags.resume, map_location='cpu')
learner_model.load_state_dict(save['model_state_dict'])
optimizer.load_state_dict(save['optimizer_state_dict'])
if flags.resume_scheduler:
scheduler.load_state_dict(save['scheduler_state_dict'])
# tune only the embedding layer
if flags.resume_strategy == 'emb':
keep = []
for group in optimizer.param_groups:
if group['params'][0].size() == (len(learner_model.vocab),
flags.demb):
keep.append(group)
optimizer.param_groups = keep
logger = logging.getLogger('logfile')
stat_keys = [
'total_loss',
'mean_episode_return',
'pg_loss',
'baseline_loss',
'entropy_loss',
'aux_loss',
'mean_win_rate',
'mean_episode_len',
]
logger.info('# Step\t%s', '\t'.join(stat_keys))
frames, stats = 0, {}
def batch_and_learn(i, lock=threading.Lock()):
"""Thread target for the learning process."""
nonlocal frames, stats
timings = prof.Timings()
while frames < flags.total_frames:
timings.reset()
batch = get_batch(free_queue, full_queue, buffers, flags, timings)
tmp_t1 = time.time()
stats = learn(model, learner_model, batch, optimizer, scheduler,
flags)
tmp_t2 = time.time()
print('learn time: ', tmp_t2 - tmp_t1)
timings.time('learn')
with lock:
to_log = dict(frames=frames)
to_log.update({k: stats[k] for k in stat_keys})
plogger.log(to_log)
frames += T * B
print('learned frames: ', frames)
if i == 0:
logging.info('Batch and learn: %s', timings.summary())
for m in range(flags.num_buffers):
free_queue.put(m)
threads = []
for i in range(flags.num_threads):
thread = threading.Thread(target=batch_and_learn,
name='batch-and-learn-%d' % i,
args=(i, ))
thread.start()
threads.append(thread)
def checkpoint():
if flags.disable_checkpoint:
return
logging.info('Saving checkpoint to %s', checkpointpath)
torch.save(
{
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'flags': vars(flags),
}, checkpointpath)
timer = timeit.default_timer
try:
last_checkpoint_time = timer()
while frames < flags.total_frames:
start_frames = frames
start_time = timer()
time.sleep(5)
if timer() - last_checkpoint_time > 10 * 60: # Save every 10 min.
checkpoint()
last_checkpoint_time = timer()
fps = (frames - start_frames) / (timer() - start_time)
if stats.get('episode_returns', None):
mean_return = 'Return per episode: %.1f. ' % stats[
'mean_episode_return']
else:
mean_return = ''
total_loss = stats.get('total_loss', float('inf'))
logging.info('After %i frames: loss %f @ %.1f fps. %sStats:\n%s',
frames, total_loss, fps, mean_return,
pprint.pformat(stats))
except KeyboardInterrupt:
return # Try joining actors then quit.
else:
for thread in threads:
thread.join()
logging.info('Learning finished after %d frames.', frames)
finally:
for _ in range(flags.num_actors):
free_queue.put(None)
for actor in actor_processes:
actor.join(timeout=1)
checkpoint()
plogger.close()
def train(flags): # pylint: disable=too-many-branches, too-many-statements
if flags.xpid is None:
flags.xpid = "torchbeast-%s" % time.strftime("%Y%m%d-%H%M%S")
plogger = file_writer.FileWriter(xpid=flags.xpid,
xp_args=flags.__dict__,
rootdir=flags.savedir)
checkpointpath = os.path.expandvars(
os.path.expanduser("%s/%s/%s" %
(flags.savedir, flags.xpid, "model.tar")))
if flags.num_buffers is None: # Set sensible default for num_buffers.
flags.num_buffers = max(2 * flags.num_actors, flags.batch_size)
if flags.num_actors >= flags.num_buffers:
raise ValueError("num_buffers should be larger than num_actors")
if flags.num_buffers < flags.batch_size:
raise ValueError("num_buffers should be larger than batch_size")
T = flags.unroll_length
B = flags.batch_size
flags.device = None
if not flags.disable_cuda and torch.cuda.is_available():
logging.info("Using CUDA.")
flags.device = torch.device("cuda")
else:
logging.info("Not using CUDA.")
flags.device = torch.device("cpu")
env = create_env(flags)
Net = get_model(flags)
model = Net(env.observation_space.shape, env.action_space, flags.use_lstm)
def get_action_shape(flags, env: gym.envs) -> typing.Tuple[int]:
# return the action shape as a tuple
if flags.env_type == 'mujoco':
return env.action_space.shape
if not flags.env_type == 'atari':
raise ValueError("Unknown env_type")
return (env.action_space.n, )
action_shape = get_action_shape(flags, env)
buffers = create_buffers(flags, env.observation_space.shape, action_shape,
get_num_agents(flags))
model.share_memory()
# Add initial RNN state.
initial_agent_state_buffers = []
for _ in range(flags.num_buffers):
state = model.initial_state(batch_size=1)
for t in state:
t.share_memory_()
initial_agent_state_buffers.append(state)
actor_processes = []
ctx = mp.get_context("fork")
free_queue = ctx.SimpleQueue()
full_queue = ctx.SimpleQueue()
for i in range(flags.num_actors):
actor = ctx.Process(
target=act,
args=(
flags,
i,
free_queue,
full_queue,
model,
buffers,
initial_agent_state_buffers,
),
)
actor.start()
actor_processes.append(actor)
Net = get_model(flags)
learner_model = Net(env.observation_space.shape, env.action_space,
flags.use_lstm).to(device=flags.device)
optimizer = torch.optim.RMSprop(
learner_model.parameters(),
lr=flags.learning_rate,
momentum=flags.momentum,
eps=flags.epsilon,
alpha=flags.alpha,
)
def lr_lambda(epoch):
return 1 - min(epoch * T * B, flags.total_steps) / flags.total_steps
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
logger = logging.getLogger("logfile")
stat_keys = [
"total_loss",
"mean_episode_return",
"pg_loss",
"baseline_loss",
"entropy_loss",
]
logger.info("# Step\t%s", "\t".join(stat_keys))
step, stats = 0, {}
def batch_and_learn(i, lock=threading.Lock()):
"""Thread target for the learning process."""
nonlocal step, stats
timings = prof.Timings()
while step < flags.total_steps:
timings.reset()
batch, agent_state = get_batch(
flags,
free_queue,
full_queue,
buffers,
initial_agent_state_buffers,
timings,
)
stats = learn(flags, model, learner_model, batch, agent_state,
optimizer, scheduler)
timings.time("learn")
with lock:
to_log = dict(step=step)
to_log.update({k: stats[k] for k in stat_keys})
plogger.log(to_log)
step += T * B
if i == 0:
logging.info("Batch and learn: %s", timings.summary())
for m in range(flags.num_buffers):
free_queue.put(m)
threads = []
for i in range(flags.num_learner_threads):
thread = threading.Thread(target=batch_and_learn,
name="batch-and-learn-%d" % i,
args=(i, ))
thread.start()
threads.append(thread)
def checkpoint():
if flags.disable_checkpoint:
return
logging.info("Saving checkpoint to %s", checkpointpath)
torch.save(
{
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(),
"flags": vars(flags),
},
checkpointpath,
)
timer = timeit.default_timer
try:
last_checkpoint_time = timer()
while step < flags.total_steps:
start_step = step
start_time = timer()
time.sleep(5)
if timer() - last_checkpoint_time > 10 * 60: # Save every 10 min.
checkpoint()
last_checkpoint_time = timer()
sps = (step - start_step) / (timer() - start_time)
if stats.get("episode_returns", None):
mean_return = ("Return per episode: %.1f. " %
stats["mean_episode_return"])
else:
mean_return = ""
total_loss = stats.get("total_loss", float("inf"))
logging.info(
"Steps %i @ %.1f SPS. Loss %f. %sStats:\n%s",
step,
sps,
total_loss,
mean_return,
pprint.pformat(stats),
)
# for mujoco, need to manually print...
print(f"sps is: {sps}")
print(pprint.pformat(stats))
except KeyboardInterrupt:
return # Try joining actors then quit.
else:
for thread in threads:
thread.join()
logging.info("Learning finished after %d steps.", step)
finally:
for _ in range(flags.num_actors):
free_queue.put(None)
for actor in actor_processes:
actor.join(timeout=1)
checkpoint()
plogger.close()
if reg_scores is not None:
reg_scores = reg_scores.view(-1, num_class, 2)
reg_scores[:, :,
0] = reg_scores[:, :, 0] * stats[1, 0] + stats[0, 0]
reg_scores[:, :,
1] = reg_scores[:, :, 1] * stats[1, 1] + stats[0, 1]
# perform stpp on scores
result_queue.put((dataset.video_list[index].id, rel_props.numpy(),
act_scores.cpu().numpy(), comp_scores.cpu().numpy(),
reg_scores.cpu().numpy(), 0))
if __name__ == '__main__':
ctx = multiprocessing.get_context(
'spawn'
) # this is crucial to using multiprocessing processes with PyTorch
# This net is used to provides setup settings. It is not used for testing.
checkpoint = torch.load(args.weights,
map_location=lambda storage, loc: storage.cuda(0))
print("model epoch {} loss: {}".format(checkpoint['epoch'],
checkpoint['best_loss']))
base_dict = {
'.'.join(k.split('.')[1:]): v
for k, v in list(checkpoint['state_dict'].items())
}
stats = checkpoint['reg_stats'].cpu().numpy()
class GPUProcess(mp.get_context("spawn").Process):
def __init__(
self,
gpu_idx: GPURank,
subprocess_init: Optional[Callable[[], None]] = None,
embedding_storage_freelist: Optional[Set[torch.FloatStorage]] = None,
) -> None:
super().__init__(daemon=True, name=f"GPU #{gpu_idx}")
self.gpu_idx = gpu_idx
self.master_endpoint, self.worker_endpoint = mp.get_context("spawn").Pipe()
self.subprocess_init = subprocess_init
self.sub_holder: Dict[
Tuple[EntityName, Partition, SubPartition],
Tuple[torch.nn.Parameter, RowAdagrad],
] = {}
self.embedding_storage_freelist = embedding_storage_freelist
@property
def my_device(self) -> torch.device:
return torch.device("cuda", index=self.gpu_idx)
def run(self) -> None:
torch.set_num_threads(1)
torch.cuda.set_device(self.my_device)
if self.subprocess_init is not None:
self.subprocess_init()
self.master_endpoint.close()
for s in self.embedding_storage_freelist:
assert s.is_shared()
cudart = torch.cuda.cudart()
res = cudart.cudaHostRegister(s.data_ptr(), s.size() * s.element_size(), 0)
torch.cuda.check_error(res)
assert s.is_pinned()
logger.info(f"GPU subprocess {self.gpu_idx} up and running")
while True:
try:
job: SubprocessArgs = self.worker_endpoint.recv()
except EOFError:
break
stats = self.do_one_job(
lhs_types=job.lhs_types,
rhs_types=job.rhs_types,
lhs_part=job.lhs_part,
rhs_part=job.rhs_part,
lhs_subpart=job.lhs_subpart,
rhs_subpart=job.rhs_subpart,
next_lhs_subpart=job.next_lhs_subpart,
next_rhs_subpart=job.next_rhs_subpart,
model=job.model,
trainer=job.trainer,
all_embs=job.all_embs,
subpart_slices=job.subpart_slices,
subbuckets=job.subbuckets,
batch_size=job.batch_size,
lr=job.lr,
)
self.worker_endpoint.send(
SubprocessReturn(gpu_idx=self.gpu_idx, stats=stats)
)
def do_one_job( # noqa
self,
lhs_types: Set[str],
rhs_types: Set[str],
lhs_part: Partition,
rhs_part: Partition,
lhs_subpart: SubPartition,
rhs_subpart: SubPartition,
next_lhs_subpart: Optional[SubPartition],
next_rhs_subpart: Optional[SubPartition],
model: MultiRelationEmbedder,
trainer: Trainer,
all_embs: Dict[Tuple[EntityName, Partition], FloatTensorType],
subpart_slices: Dict[Tuple[EntityName, Partition, SubPartition], slice],
subbuckets: Dict[
Tuple[int, int], Tuple[LongTensorType, LongTensorType, LongTensorType]
],
batch_size: int,
lr: float,
) -> Stats:
tk = TimeKeeper()
for embeddings in all_embs.values():
assert embeddings.is_pinned()
occurrences: Dict[
Tuple[EntityName, Partition, SubPartition], Set[Side]
] = defaultdict(set)
for entity_name in lhs_types:
occurrences[entity_name, lhs_part, lhs_subpart].add(Side.LHS)
for entity_name in rhs_types:
occurrences[entity_name, rhs_part, rhs_subpart].add(Side.RHS)
if lhs_part != rhs_part: # Bipartite
assert all(len(v) == 1 for v in occurrences.values())
tk.start("copy_to_device")
for entity_name, part, subpart in occurrences.keys():
if (entity_name, part, subpart) in self.sub_holder:
continue
embeddings = all_embs[entity_name, part]
optimizer = trainer.partitioned_optimizers[entity_name, part]
subpart_slice = subpart_slices[entity_name, part, subpart]
# TODO have two permanent storages on GPU and move stuff in and out
# from them
# logger.info(f"GPU #{self.gpu_idx} allocating {(subpart_slice.stop - subpart_slice.start) * embeddings.shape[1] * 4:,} bytes")
gpu_embeddings = torch.empty(
(subpart_slice.stop - subpart_slice.start, embeddings.shape[1]),
dtype=torch.float32,
device=self.my_device,
)
gpu_embeddings.copy_(embeddings[subpart_slice], non_blocking=True)
gpu_embeddings = torch.nn.Parameter(gpu_embeddings)
gpu_optimizer = RowAdagrad([gpu_embeddings], lr=lr)
(cpu_state,) = optimizer.state.values()
(gpu_state,) = gpu_optimizer.state.values()
# logger.info(f"GPU #{self.gpu_idx} allocating {(subpart_slice.stop - subpart_slice.start) * 4:,} bytes")
gpu_state["sum"].copy_(cpu_state["sum"][subpart_slice], non_blocking=True)
self.sub_holder[entity_name, part, subpart] = (
gpu_embeddings,
gpu_optimizer,
)
logger.debug(
f"Time spent copying subparts to GPU: {tk.stop('copy_to_device'):.4f} s"
)
for (
(entity_name, part, subpart),
(gpu_embeddings, gpu_optimizer),
) in self.sub_holder.items():
for side in occurrences[entity_name, part, subpart]:
model.set_embeddings(entity_name, side, gpu_embeddings)
trainer.partitioned_optimizers[
entity_name, part, subpart
] = gpu_optimizer
tk.start("translate_edges")
num_edges = subbuckets[lhs_subpart, rhs_subpart][0].shape[0]
edge_perm = torch.randperm(num_edges)
edges_lhs, edges_rhs, edges_rel = subbuckets[lhs_subpart, rhs_subpart]
_C.shuffle(edges_lhs, edge_perm, os.cpu_count())
_C.shuffle(edges_rhs, edge_perm, os.cpu_count())
_C.shuffle(edges_rel, edge_perm, os.cpu_count())
assert edges_lhs.is_pinned()
assert edges_rhs.is_pinned()
assert edges_rel.is_pinned()
gpu_edges = EdgeList(
EntityList.from_tensor(edges_lhs),
EntityList.from_tensor(edges_rhs),
edges_rel,
).to(self.my_device, non_blocking=True)
logger.debug(f"GPU #{self.gpu_idx} got {num_edges} edges")
logger.debug(
f"Time spent copying edges to GPU: {tk.stop('translate_edges'):.4f} s"
)
tk.start("processing")
stats = process_in_batches(
batch_size=batch_size, model=model, batch_processor=trainer, edges=gpu_edges
)
logger.debug(f"Time spent processing: {tk.stop('processing'):.4f} s")
next_occurrences: Dict[
Tuple[EntityName, Partition, SubPartition], Set[Side]
] = defaultdict(set)
if next_lhs_subpart is not None:
for entity_name in lhs_types:
next_occurrences[entity_name, lhs_part, next_lhs_subpart].add(Side.LHS)
if next_rhs_subpart is not None:
for entity_name in rhs_types:
next_occurrences[entity_name, rhs_part, next_rhs_subpart].add(Side.RHS)
tk.start("copy_from_device")
for (entity_name, part, subpart), (gpu_embeddings, gpu_optimizer) in list(
self.sub_holder.items()
):
if (entity_name, part, subpart) in next_occurrences:
continue
embeddings = all_embs[entity_name, part]
optimizer = trainer.partitioned_optimizers[entity_name, part]
subpart_slice = subpart_slices[entity_name, part, subpart]
embeddings[subpart_slice].data.copy_(
gpu_embeddings.detach(), non_blocking=True
)
del gpu_embeddings
(cpu_state,) = optimizer.state.values()
(gpu_state,) = gpu_optimizer.state.values()
cpu_state["sum"][subpart_slice].copy_(gpu_state["sum"], non_blocking=True)
del gpu_state["sum"]
del self.sub_holder[entity_name, part, subpart]
logger.debug(
f"Time spent copying subparts from GPU: {tk.stop('copy_from_device'):.4f} s"
)
logger.debug(f"do_one_job: Time unaccounted for: {tk.unaccounted():.4f} s")
return stats
from .zmq_adapter import InitSender, WaitAll, SendAll
import abc
import torch.multiprocessing as _mp
mp = _mp.get_context('spawn')
class Simulator(mp.Process):
'''
Wrapper for simulator.
Functions to override:
on_init: Initialization after the process has started.
restart: restart the environment.
terminal: property that tell whether the game has reached terminal
get_key: from the key, get the content. e.g. ``get_key("s")`` will give the encoded state of the game.
set_key: set the key from replies. e.g., ``set_key("a", 2)`` set the action to be 2 (and the underlying game can continue).
'''
def __init__(self, id, desc):
'''
Example:
desc = dict(
actor = dict(
input = dict(s="", last_terminal=""),
reply = dict(a="")
connector = "name1"
),
train = dict(
input = dict(s="", r="", a=""),
reply = None,
connector = "name2"
)
}
'''
def test_cuda(self):
torch.cuda.FloatTensor([1]) # initialize CUDA outside of leak checker
self._test_sharing(mp.get_context('spawn'), torch.cuda.FloatTensor)
def __init__(
self,
make_env_fn: Callable[..., Union[Env, RLEnv]] = _make_env_fn,
env_fn_args: Sequence[Tuple] = None,
auto_reset_done: bool = True,
multiprocessing_start_method: str = "forkserver",
workers_ignore_signals: bool = False,
) -> None:
"""..
:param make_env_fn: function which creates a single environment. An
environment can be of type :ref:`env.Env` or :ref:`env.RLEnv`
:param env_fn_args: tuple of tuple of args to pass to the
:ref:`_make_env_fn`.
:param auto_reset_done: automatically reset the environment when
done. This functionality is provided for seamless training
of vectorized environments.
:param multiprocessing_start_method: the multiprocessing method used to
spawn worker processes. Valid methods are
:py:`{'spawn', 'forkserver', 'fork'}`; :py:`'forkserver'` is the
recommended method as it works well with CUDA. If :py:`'fork'` is
used, the subproccess must be started before any other GPU useage.
:param workers_ignore_signals: Whether or not workers will ignore SIGINT and SIGTERM
and instead will only exit when :ref:`close` is called
"""
self._is_waiting = False
self._is_closed = True
assert (
env_fn_args is not None and len(env_fn_args) > 0
), "number of environments to be created should be greater than 0"
self._num_envs = len(env_fn_args)
assert multiprocessing_start_method in self._valid_start_methods, (
"multiprocessing_start_method must be one of {}. Got '{}'").format(
self._valid_start_methods, multiprocessing_start_method)
self._auto_reset_done = auto_reset_done
self._mp_ctx = mp.get_context(multiprocessing_start_method)
self._workers = []
(
self._connection_read_fns,
self._connection_write_fns,
) = self._spawn_workers( # noqa
env_fn_args,
make_env_fn,
workers_ignore_signals=workers_ignore_signals,
)
self._is_closed = False
for write_fn in self._connection_write_fns:
write_fn((OBSERVATION_SPACE_COMMAND, None))
self.observation_spaces = [
read_fn() for read_fn in self._connection_read_fns
]
for write_fn in self._connection_write_fns:
write_fn((ACTION_SPACE_COMMAND, None))
self.action_spaces = [
read_fn() for read_fn in self._connection_read_fns
]
for write_fn in self._connection_write_fns:
write_fn((NUMBER_OF_EPISODES_COMMAND, None))
self.number_of_episodes = [
read_fn() for read_fn in self._connection_read_fns
]
self._paused: List[Tuple] = []
if __name__ == '__main__':
parser = get_parser()
argv = sys.argv[1:]
args, _ = parser.parse_known_args(argv)
if args.config == 'multiwoz':
config = MultiWozConfig()
else:
raise NotImplementedError(
'Config of the dataset {} not implemented'.format(args.config))
init_logging_handler(args.log_dir)
logging.debug(str(args))
try:
mp = mp.get_context('spawn')
except RuntimeError:
pass
if args.pretrain:
pass
elif args.test:
pass
else: # training
current_time = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
logging.debug('train {}'.format(current_time))
print("training start")
output[cnt: cnt + sc.size(0), :] = sc
cnt += sc.size(0)
act_scores, comp_scores, reg_scores = reorg_stpp.forward(output, prop_ticks, prop_scaling)
if reg_scores is not None:
reg_scores = reg_scores.view(-1, num_class, 2)
reg_scores[:, :, 0] = reg_scores[:, :, 0] * stats[1, 0] + stats[0, 0]
reg_scores[:, :, 1] = reg_scores[:, :, 1] * stats[1, 1] + stats[0, 1]
# perform stpp on scores
result_queue.put((dataset.video_list[index].id, rel_props.numpy(), act_scores.cpu().numpy(), \
comp_scores.cpu().numpy(), reg_scores.cpu().numpy(), output.cpu().numpy()))
if __name__ == '__main__':
ctx = multiprocessing.get_context('spawn') # this is crucial to using multiprocessing processes with PyTorch
# This net is used to provides setup settings. It is not used for testing.
net = SSN(num_class, 2, 5, 2,
args.modality, test_mode=True,
base_model=args.arch, no_regression=args.no_regression, stpp_cfg=stpp_configs)
if args.test_crops == 1:
cropping = torchvision.transforms.Compose([
GroupScale(net.scale_size),
GroupCenterCrop(net.input_size),
])
elif args.test_crops == 10:
cropping = torchvision.transforms.Compose([
GroupOverSample(net.input_size, net.scale_size)
])
def test_cuda_variable_sharing(self):
for requires_grad in [True, False]:
var = torch.arange(1., 26, device='cuda').view(5, 5).requires_grad_(requires_grad)
self._test_autograd_sharing(var, mp.get_context('spawn'))
