moved the order of  moving to device and creating the optimizer
@@ -102,9 +102,28 @@ class Transformer:
verbose=True,
seed=None,
):
-
+ # get device
device, num_gpus = get_device(num_gpus=n_gpu, local_rank=-1)
+ # unwrap model
+ if isinstance(self.model, torch.nn.DataParallel):
+ self.model = self.model.module
+
+ # wrap in DataParallel or DistributedDataParallel
+ if local_rank != -1:
+ self.model = torch.nn.parallel.DistributedDataParallel(
+ self.model,
+ device_ids=[local_rank],
+ output_device=local_rank,
+ find_unused_parameters=True,
+ )
+ else:
+ if num_gpus > 1:
+ self.model = torch.nn.DataParallel(self.model, device_ids=list(range(num_gpus)))
+
+ # move to device
+ self.model.to(device)
+
if seed is not None:
Transformer.set_seed(seed, num_gpus > 0)
@@ -116,6 +135,7 @@ class Transformer:
else:
t_total = len(train_dataloader) // gradient_accumulation_steps * num_train_epochs
+ # set optimizer
if optimizer is None:
no_decay = ["bias", "LayerNorm.weight"]
optimizer_grouped_parameters = [
@@ -138,6 +158,7 @@ class Transformer:
]
optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate, eps=adam_epsilon)
+ # set scheduler
if scheduler is None:
scheduler = get_linear_schedule_with_warmup(
optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total
@@ -150,30 +171,16 @@ class Transformer:
raise ImportError("Please install apex from https://www.github.com/nvidia/apex")
self.model, optimizer = amp.initialize(self.model, optimizer, opt_level=fp16_opt_level)
- if local_rank != -1:
- self.model = torch.nn.parallel.DistributedDataParallel(
- self.model,
- device_ids=[local_rank],
- output_device=local_rank,
- find_unused_parameters=True,
- )
- else:
- if isinstance(self.model, torch.nn.DataParallel):
- self.model = self.model.module
-
- if num_gpus > 1:
- self.model = torch.nn.DataParallel(self.model, device_ids=list(range(num_gpus)))
-
- self.model.to(device)
- self.model.train()
-
+ # init training
global_step = 0
tr_loss = 0.0
+ self.model.train()
self.model.zero_grad()
train_iterator = trange(
int(num_train_epochs), desc="Epoch", disable=local_rank not in [-1, 0] or not verbose
)
+ # train
for _ in train_iterator:
epoch_iterator = tqdm(
train_dataloader, desc="Iteration", disable=local_rank not in [-1, 0] or not verbose
@@ -214,9 +221,6 @@ class Transformer:
train_iterator.close()
break
- # empty cache
- del [batch]
- torch.cuda.empty_cache()
return global_step, tr_loss / global_step
def predict(self, eval_dataloader, get_inputs, n_gpu=1, verbose=True):
