I’ve tried everything. gc.collect, torch.cuda.empty_cache, deleting every possible tensor and variable as soon as it is used, setting batch size to 1, nothing seems to work.
I’m working on text to code generation problem and utilizing the code from this repository : TranX
I’ve rewritten the data loader, model training pipeline and have made it as simple as i possibly can, but somehow it always runs out of memory.
Here’s my data loader :
class DSET:
def __init__(self, load_file='', mode='train', batch_size=16):
self.load_file = load_file
self.data = open(self.load_file, 'r', encoding='utf8').readlines()
class EXAMPLE(object):
def __init__(self, examples):
self._examples = examples
self.generator = AST('python')
def generate_ast(self, code_text):
tree = self.generator.generate_ast(code_text)
return tree
def get_actions(self, ast_tree):
actions = self.generator.get_actions(ast_tree)
return actions
def __getitem__(self, index):
line = json.loads(self._examples[index])
try:
canonical_intent, slot_map = canonicalize_intent(
line['nl'])
canonical_code = canonicalize_code(line['code'], slot_map)
intent_tokens = tokenize_intent(canonical_intent)
python_ast = ast.parse(canonical_code)
canonical_code_source = astor.to_source(python_ast).strip()
target_ast = self.generate_ast(python_ast)
target_actions = self.get_actions(target_ast)
target_action_infos = get_action_infos(
intent_tokens, target_actions)
example = Example(src_sent=intent_tokens,
tgt_actions=target_action_infos,
tgt_code=canonical_code_source,
tgt_ast=target_ast,
meta=dict(example_dict=line,
slot_map=slot_map))
except Exception as e:
example = Example(src_sent=[''],
tgt_actions=[None],
tgt_code=[''],
tgt_ast=None,
meta={})
return example
def __setitem__(self, index, value):
pass
self.examples = EXAMPLE(self.data)
self.num_examples = len(self.data)
self.example_indices = np.arange(self.num_examples)
self.iterator = 0
self.batch_size = batch_size
self.num_steps = self.num_examples // self.batch_size
def batch_iter(self):
indices = self.example_indices[self.iterator *
self.batch_size:(self.iterator + 1) *
self.batch_size]
x = []
for idx in indices:
sample = self.examples[idx]
x.append(sample)
self.iterator += 1
if self.iterator > len(self.example_indices) // self.batch_size:
self.iterator = 0
np.random.shuffle(self.example_indices)
x.sort(key=lambda e: -len(e.src_sent))
return x
def __len__(self):
return self.num_examples
Here’s the training code :
def train(args):
train_set = DSET(args.train_file, mode='train', batch_size=args.batch_size)
validation_set = DSET(args.dev_file,
mode='valid',
batch_size=args.batch_size)
vocab = pickle.load(open(args.vocab, 'rb'))
grammar = ASDLGrammar.from_text(open(args.asdl_file).read())
transition_system = Registrable.by_name(args.transition_system)(grammar)
parser = Registrable.by_name(args.parser)
model = parser(args, vocab, transition_system)
model.train()
evaluator = Registrable.by_name(args.evaluator)(transition_system,
args=args)
if args.cuda:
model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
nn_utils.glorot_init(model.parameters())
print('Begining Training ...')
print(
f'Training Examples : {len(train_set)}, Validation Examples : {len(validation_set)}'
)
print(
f'Number Of Epochs : {args.max_epoch}, Steps Per Epoch : {train_set.num_steps}, Batch Size : {args.batch_size}'
)
total_steps = 0
for e in range(args.max_epoch):
start_time = time.time()
for iter in range(train_set.num_steps):
total_steps += 1
loss_val = 0
batch = train_set.batch_iter()
optimizer.zero_grad()
loss = -model.score(batch)
loss = torch.mean(loss)
if iter % args.log_every == 0:
loss_val = torch.sum(loss).detach().item()
print(f'Iteration : {iter}, Loss : {loss_val}')
loss.backward(retain_graph=False)
optimizer.step()
del loss
del loss_val
del batch
gc.collect()
torch.cuda.empty_cache()
time.sleep(2)
print(f'Epoch {e} : Time Taken : {time.time() - start_time}')
model_file = args.save_to + '/' + f'model.epoch_{e}.bin'
print(f'Saved Model To : {model_file}')
model.save(model_file)
torch.save(optimizer.state_dict(),
args.save_to + '/' + f'model_{e}.optim.bin')
# Run Validation
print('Running Validation...')
try:
eval_results = evaluate(validation_set.examples,
model,
evaluator,
args,
verbose=False,
eval_top_pred_only=args.eval_top_pred_only)
validation_score = eval_results[evaluator.default_metric]
print(f'Evaluation Score: {validation_score}')
except Exception as e:
print(f'Could not validate: {e}')
pass
The model.score method is custom by the repo author and i’ve added delete and gc.collect and torch.cuda.empty_cache lines throughout
I’m running pytorch 1.9.1 with cuda 11.1 on a 16gb GPU instance on aws ec2 with 32gb ram and ubuntu 18.04
I’ve re-written the code to make it more efficient as the code in the repository loaded the whole bin file of the dataset at once.
But i can’t train the model, even with batch size of 1.
With batch size of 8, it crashes after 46 iterations, with batch size of 1 it goes upto 48k iterations but then crashes.
Traceback for GPU (After 46 Iterations):
Traceback (most recent call last):
File "train.py", line 150, in <module>
train(args)
File "train.py", line 98, in train
loss.backward(retain_graph=False)
File "/usr/local/lib/python3.6/dist-packages/comet_ml/monkey_patching.py", lin e 312, in wrapper
return_value = original(*args, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/torch/_tensor.py", line 255, in b ackward
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=i nputs)
File "/usr/local/lib/python3.6/dist-packages/torch/autograd/__init__.py", line 149, in backward
allow_unreachable=True, accumulate_grad=True) # allow_unreachable flag
RuntimeError: CUDA out of memory. Tried to allocate 7.27 GiB (GPU 0; 14.76 GiB t
Traceback for CPU (It simply says Killed) but goes upto iteration 672:
Iteration : 672, Loss : 702.1222534179688
Killed
Running dmesg gives the following output
[2059991.491436] oom-kill:constraint=CONSTRAINT_NONE,nodemask=(null),cpuset=/,me ms_allowed=0,global_oom,task_memcg=/user.slice,task=python3,pid=25315,uid=0
[2059991.491542] Out of memory: Killed process 25315 (python3) total-vm:53312244 kB, anon-rss:31451456kB, file-rss:74816kB, shmem-rss:12296kB, UID:0 pgtables:690 68kB oom_score_adj:0
[2059992.056260] oom_reaper: reaped process 25315 (python3), now anon-rss:0kB, f