I've trained a model with a single process. It worked fine. Now I want to train again but with MultipleProcess and I'm not sure where/how I should save my model's weights when I detect the best loss. For the 'how', I'm following this guide:
https://github.com/pytorch/xla/blob/master/API_GUIDE.md#saving-and-loading-xla-tensors
But in the loop/spawn where should I add it? Here is an extract (simplified of my code). I want to track the best_valid_logloss and then save weights the matching epoch.
def run_stage(index, flags):
torch.manual_seed(seed)
transform_fwd = transforms.Compose([transforms.ToTensor(),
transforms.Normalize(IMG_MEAN, IMG_STD)])
l_device = xm.xla_device()
world_size = xm.xrt_world_size()
X_train, X_valid = flags["X_train"], flags["X_valid"]
num_workers, log_steps = flags["num_workers"], flags["log_steps"]
# Scale learning rate to num cores
BASE_LR, MIN_LR = flags["BASE_LR"] * world_size, flags["MIN_LR"] * world_size
dataset_train = MyCustomDataset(X_train, subset='train', transform=transform_fwd, augmentations=aug)
dataset_valid = MyCustomDataset(X_valid, subset='valid', transform=transform_fwd)
# Creates the (distributed) train sampler, which let this process only access its portion of the training dataset.
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset_train, num_replicas=world_size, rank=xm.get_ordinal(), shuffle=True)
valid_sampler = torch.utils.data.distributed.DistributedSampler(dataset_valid, num_replicas=world_size, rank=xm.get_ordinal(), shuffle=False)
# Creates dataloaders, which load data in batches
train_batches = DataLoader(dataset_train, sampler=train_sampler, batch_size=BATCH_SIZE, num_workers=num_workers, drop_last=True)
valid_batches = DataLoader(dataset_valid, sampler=valid_sampler, batch_size=BATCH_SIZE, shuffle=False, num_workers=num_workers, drop_last=True)
# Note: each process has its own identical copy of the model
cnn_model, loss, optimizer = build_model(BASE_LR, dev=l_device)
def train_loop_fn(train_batches):
cnn_model.train()
tracker = xm.RateTracker()
for x, train_batch in enumerate(train_batches):
try:
...
optimizer.zero_grad()
img_data = img_data.to(l_device)
labels_data = labels_data.to(l_device) # (BS, 1)
input_v = Variable(img_data)
output = cnn_model(input_v.float())
loss_dis = loss(output, Variable(labels_data, requires_grad=False))
loss_dis.backward() # backward pass
xm.optimizer_step(optimizer)
...
tracker.add(BATCH_SIZE)
if x % log_steps == 0:
print('[xla:{}]({}) Loss={:.5f} Rate={:.2f} GlobalRate={:.2f} Time={}'.format(
xm.get_ordinal(), x, loss_dis.item(), tracker.rate(),
tracker.global_rate(), time.asctime()), flush=True)
except Exception as ex:
print("Training batch error:", ex)
# Compute metrics
...
print('[xla:{}] Train LogLoss={:.5f}, Train Accuracy={:.2f}%'.format(
xm.get_ordinal(), ll_train, acc_train), flush=True)
return (loss_train, acc_train, ll_train)
def valid_loop_fn(valid_batches):
cnn_model.eval()
for valid_batch in valid_batches:
try:
img_data = img_data.to(l_device)
labels_data = labels_data.to(l_device)
input_v = Variable(img_data)
with torch.no_grad():
output = cnn_model(input_v.float())
# Compute loss ...
...
except Exception as ex:
print("Validation batch error:", ex)
...
# Compute metrics
print('[xla:{}] Valid LogLoss={:.5f}, Valid Accuracy={:.2f}%'.format(
xm.get_ordinal(), ll_test, acc_test), flush=True)
return (loss_test, acc_test, ll_test)
text_writer = None
if xm.is_master_ordinal():
text_writer = open(report_path, 'a' if opt.resume > 0 else 'w')
history = []
best_valid_logloss = 99999.0
for epoch in range(1,20):
lr = BASE_LR
###########
# Training
para_loader = pl.ParallelLoader(train_batches, [l_device])
loss_train, acc_train, ll_train = train_loop_fn(para_loader.per_device_loader(l_device))
xm.master_print("Finished training epoch {}".format(epoch))
#############
# Validation
para_loader = pl.ParallelLoader(valid_batches, [l_device])
loss_test, acc_test, ll_test = valid_loop_fn(para_loader.per_device_loader(l_device))
print("\n[%s]"%index, '[Epoch %d] Train loss: %.4f acc: %.3f logloss: %.4f | Valid loss: %.4f acc: %.3f logloss: %.4f | LR: %.6f' % (epoch, loss_train, acc_train, ll_train, loss_test, acc_test, ll_test, lr))
history.append((epoch, lr, loss_train, acc_train, ll_train, loss_test, acc_test, ll_test))
if text_writer is not None: text_writer.write('%d,%.4f,%.4f,%.3f,%.4f,%.4f,%.3f,%.4f\n' % (epoch, lr, loss_train, acc_train, ll_train, loss_test, acc_test, ll_test))
if text_writer is not None: text_writer.flush()
if ll_test < best_valid_logloss:
print("[%s]"%index, '[Epoch %d] Valid logloss improved from %.4f to %.4f ... Saving model' % (epoch, best_valid_logloss, ll_test))
best_valid_logloss = ll_test
xm.save(cnn_model.state_dict(), os.path.join(snapshot_path, 'model_best.pt'))
if text_writer is not None: text_writer.close()
And my spawn call:
xmp.spawn(run_stage, args=(FLAGS,), nprocs=FLAGS["num_workers"], start_method='fork')
Do you have any guide or example to perform it correctly?
LSOFT123
All 8 comments
Hello,
Using xm.save to save your model weights (or any other xla tensor) is the recommended way.
Saving when the "new best loss" is detected is fine, but we should be sure that every core detects the best loss at the same time.
The underlying rule is; every core should do the same tensor work. Example, if the code shards the validation data, and computed validation losses are different from one device to the next, then sometimes one device will get to the line with xm.save and the next device won't get there, and that can cause a problem. This issue will not happen if every core does the same validation work (validates the whole validation set for example).
Taking a quick look at the code above, there may be a subtle version of the same issue in this line:
if text_writer is not None: text_writer.write('%d,%.4f,%.4f,%.3f,%.4f,%.4f,%.3f,%.4f\n' % (epoch, lr, loss_train, acc_train, ll_train, loss_test, acc_test, ll_test))
if the loss values (e.g. loss_test) are tensors, then this line will send the tensors to cpu and write to disk, only on the master device. This will not happen on the other devices, which violates the principle above.
if all of what's being printed are already on cpu (sent to cpu on all devices before this line), then there's no issue.
taylanbil
on 19 Mar 2020
Thanks a lot for your help. So in my code I should add another step for the whole validation data (without valid_sampler) like below. Does it look correct to you?
def run_stage(index, flags):
...
# Creates dataloaders, which load data in batches
train_batches = DataLoader(dataset_train, sampler=train_sampler, batch_size=BATCH_SIZE, num_workers=num_workers, drop_last=True)
valid_batches = DataLoader(dataset_valid, sampler=valid_sampler, batch_size=BATCH_SIZE, shuffle=False, num_workers=num_workers, drop_last=True)
...
# Full validation without sampler so all cores have the same
full_valid_batches = DataLoader(dataset_valid, batch_size=BATCH_SIZE, shuffle=False, num_workers=num_workers, drop_last=True)
...
history = []
best_valid_logloss = 99999.0
for epoch in range(1,20):
##############
# Training (with train_sampler)
para_loader = pl.ParallelLoader(train_batches, [l_device])
loss_train, acc_train, ll_train = train_loop_fn(para_loader.per_device_loader(l_device))
xm.master_print("Finished training epoch {}".format(epoch))
###############
# Validation (with valid_sampler)
para_loader = pl.ParallelLoader(valid_batches, [l_device])
valid_loop_fn(para_loader.per_device_loader(l_device))
###############
# Full Validation (without valid_sampler)
para_loader = pl.ParallelLoader(full_valid_batches, [l_device])
loss_test, acc_test, ll_test = valid_loop_fn(para_loader.per_device_loader(l_device))
if ll_test < best_valid_logloss:
print("[%s]"%index, '[Epoch %d] Valid logloss improved from %.4f to %.4f ... Saving model' % (epoch, best_valid_logloss, ll_test))
best_valid_logloss = ll_test
xm.save(cnn_model.state_dict(), os.path.join(snapshot_path, 'model_best.pt'))
xmp.spawn(run_stage, args=(FLAGS,), nprocs=FLAGS["num_workers"], start_method='fork')
This should be fine, but I would just do full validation, unless you have a specific reason to do 2 different validation loops every time.
(quick tip: the print statement towards the end will print num_cores many times, you could use xm.master_print to print only from master device).
taylanbil
on 19 Mar 2020
Thanks. Looks good now. Training is super fast with 8 cores. I've kept only one full validation.
2 more questions about learning rate with MP if you don't mind:
Learning rate must be always scaled, correct?
With:
learning_rate = FLAGS['learning_rate'] * xm.xrt_world_size()
If learning_rate = 0.01 for single core, then it will become 0.08 for 8 cores.
The rational behind is the subset of data (train_sampler)? But such scaling does not guarantee that we would get same results as single core, correct? And what would be wrong to not scale it?
Also, how to setup correctly a scheduler on learning rate?
There is a sample here:
https://github.com/pytorch/xla/blob/master/test/test_train_mp_imagenet.py
With a learning rate wrapper but where can we pass how own scheduler (such as StepLR):
lr_scheduler = schedulers.wrap_optimizer_with_scheduler(
optimizer,
scheduler_type=getattr(FLAGS, 'lr_scheduler_type', None),
scheduler_divisor=getattr(FLAGS, 'lr_scheduler_divisor', None),
scheduler_divide_every_n_epochs=getattr(
FLAGS, 'lr_scheduler_divide_every_n_epochs', None),
num_steps_per_epoch=num_training_steps_per_epoch,
summary_writer=writer)
Finally, test_utils looks great, is there any documentation on it?
...
writer = test_utils.get_summary_writer(FLAGS.logdir)
...
test_utils.print_test_update(device, accuracy)
...
test_utils.write_to_summary(writer, epoch,
dict_to_write={'Accuracy/test': accuracy},
write_xla_metrics=True)
...
test_utils.close_summary_writer(writer)
...
If you wanted a new custom LR scheduler, you can add it that file as a subclass of _LRScheduler like this example and then add it as an option in the wrap_optimizer_with_scheduler method like here.
Feel free to send us a PR with that change, we'd love to have more schedulers if you think it's a good one!
For test_utils, all the methods have docstrings as documentation. Let me know if you had any further questions about any of those methods
zcain117
on 20 Mar 2020
Thanks.
The MP training (described in previous messages in this thread) crashes after a few epoch with this error.
Exception in device=TPU:6: Aborted: Session 447f9a616c838ab9 is not found.
Exception in thread Thread-10:
Traceback (most recent call last):
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/threading.py", line 916, in _bootstrap_inner
self.run()
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/threading.py", line 864, in run
self._target(*self._args, **self._kwargs)
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/parallel_loader.py", line 165, in _worker
batch = xm.send_cpu_data_to_device(batch, device)
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 518, in send_cpu_data_to_device
return ToXlaTensorArena(convert_fn, select_fn).transform(data)
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 290, in transform
self._convert()
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 262, in _convert
self._converted_tensors = self._convert_fn(self._tensors)
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 513, in convert_fn
return torch_xla._XLAC._xla_tensors_from_aten(tensors, devices)
RuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:367 : Check failed: session->session()->Run( session_work->feed_inputs, session_work->outputs_handles, &outputs) == ::tensorflow::Status::OK() (Aborted: Session 34154edd62a567a2 is not found. vs. OK)
*** Begin stack trace ***
tensorflow::CurrentStackTrace[abi:cxx11]()
clone
*** End stack trace ***
Traceback (most recent call last):
Exception in thread Thread-10:
Traceback (most recent call last):
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/threading.py", line 916, in _bootstrap_inner
self.run()
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/threading.py", line 864, in run
self._target(*self._args, **self._kwargs)
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/parallel_loader.py", line 165, in _worker
batch = xm.send_cpu_data_to_device(batch, device)
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 518, in send_cpu_data_to_device
return ToXlaTensorArena(convert_fn, select_fn).transform(data)
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 290, in transform
self._convert()
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 262, in _convert
self._converted_tensors = self._convert_fn(self._tensors)
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 513, in convert_fn
return torch_xla._XLAC._xla_tensors_from_aten(tensors, devices)
RuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:367 : Check failed: session->session()->Run( session_work->feed_inputs, session_work->outputs_handles, &outputs) == ::tensorflow::Status::OK() (Aborted: Session e6b49aa4dad53fb6 is not found. vs. OK)
*** Begin stack trace ***
tensorflow::CurrentStackTrace[abi:cxx11]()
clone
*** End stack trace ***
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 119, in _start_fn
fn(gindex, *args)
File "<ipython-input-55-0466645c4f1b>", line 231, in run_stage
loss_train, acc_train, ll_train = train_loop_fn(para_loader.per_device_loader(l_device))
File "<ipython-input-55-0466645c4f1b>", line 78, in train_loop_fn
for x, train_batch in enumerate(train_batches):
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/parallel_loader.py", line 31, in __next__
return self.next()
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/parallel_loader.py", line 34, in next
xm.mark_step()
File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 405, in mark_step
wait=xu.getenv_as('XLA_SYNC_WAIT', bool, False))
RuntimeError: Aborted: Session 447f9a616c838ab9 is not found.
---------------------------------------------------------------------------
Exception Traceback (most recent call last)
<ipython-input-59-f1c236111cf9> in <module>
31 "num_workers": 8, "log_steps": 500
32 }
---> 33 xmp.spawn(run_stage, args=(FLAGS,), nprocs=FLAGS["num_workers"], start_method='fork') # m1, history1
34
35 # plot_history(history1)
/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py in spawn(fn, args, nprocs, join, daemon, start_method)
180 join=join,
181 daemon=daemon,
--> 182 start_method=start_method)
/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py in start_processes(fn, args, nprocs, join, daemon, start_method)
156
157 # Loop on join until it returns True or raises an exception.
--> 158 while not context.join():
159 pass
160
/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py in join(self, timeout)
106 raise Exception(
107 "process %d terminated with signal %s" %
--> 108 (error_index, name)
109 )
110 else:
Exception: process 6 terminated with signal SIGABRT
re: scaling the learning rate by world size
There is no rule that governs the learning rate, it's a tunable parameter as always.
If I understand your question correctly, I think you are trying to establish an equivalence b/w training on 1 core and 8 (more generally, N) cores. I don't believe you can do that w/ tuning the learning rate, but rather the batch size.
- After every training step, we average gradients here.
- So when we take an optimizer step, gradients have been reduced.
- Thus, training under the following scenarios is equivalent (assuming gradients are additive):
- 8 cores w/ LR=0.1 and batch size per core 128
- 1 core w/ LR=0.1 and batch size 128*8=1024
taylanbil
on 20 Mar 2020
re: crash above
This seems like your tpu went away (maybe hit a maintenance event, or was preempted?). does this happen again if you run again? If so, it may be a bug and let's track it in a separate issue. Please let us know if you have more questions about the original subject of this issue.
taylanbil
on 20 Mar 2020
Related issues
ailzhang
路
6Comments
ibeltagy
路
5Comments
mruberry
路
4Comments
patrickvonplaten
路
7Comments
myleott
路
6Comments