178 lines
5.8 KiB
Python
178 lines
5.8 KiB
Python
# Copyright The PyTorch Lightning team.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import gc
|
|
import time
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from tqdm import tqdm
|
|
|
|
from pytorch_lightning import LightningModule, seed_everything, Trainer
|
|
from tests.base.models import ParityModuleMNIST, ParityModuleRNN
|
|
|
|
|
|
def assert_parity_relative(pl_values, pt_values, norm_by: float = 1, max_diff: float = 0.1):
|
|
# assert speeds
|
|
diffs = np.asarray(pl_values) - np.mean(pt_values)
|
|
# norm by vanilla time
|
|
diffs = diffs / norm_by
|
|
# relative to mean reference value
|
|
diffs = diffs / np.mean(pt_values)
|
|
assert np.mean(diffs) < max_diff, f"Lightning diff {diffs} was worse than vanilla PT (threshold {max_diff})"
|
|
|
|
|
|
def assert_parity_absolute(pl_values, pt_values, norm_by: float = 1, max_diff: float = 0.55):
|
|
# assert speeds
|
|
diffs = np.asarray(pl_values) - np.mean(pt_values)
|
|
# norm by event count
|
|
diffs = diffs / norm_by
|
|
assert np.mean(diffs) < max_diff, f"Lightning {diffs} was worse than vanilla PT (threshold {max_diff})"
|
|
|
|
|
|
# ParityModuleMNIST runs with num_workers=1
|
|
@pytest.mark.parametrize(
|
|
'cls_model,max_diff_speed,max_diff_memory',
|
|
[
|
|
(ParityModuleRNN, 0.05, 0.0),
|
|
(ParityModuleMNIST, 0.25, 0.0), # todo: lower this thr
|
|
]
|
|
)
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
|
|
def test_pytorch_parity(
|
|
tmpdir,
|
|
cls_model: LightningModule,
|
|
max_diff_speed: float,
|
|
max_diff_memory: float,
|
|
num_epochs: int = 4,
|
|
num_runs: int = 3,
|
|
):
|
|
"""
|
|
Verify that the same pytorch and lightning models achieve the same results
|
|
"""
|
|
lightning = measure_loops(cls_model, kind="PT Lightning", num_epochs=num_epochs, num_runs=num_runs)
|
|
vanilla = measure_loops(cls_model, kind="Vanilla PT", num_epochs=num_epochs, num_runs=num_runs)
|
|
|
|
# make sure the losses match exactly to 5 decimal places
|
|
print(f"Losses are for... \n vanilla: {vanilla['losses']} \n lightning: {lightning['losses']}")
|
|
for pl_out, pt_out in zip(lightning['losses'], vanilla['losses']):
|
|
np.testing.assert_almost_equal(pl_out, pt_out, 5)
|
|
|
|
# drop the first run for initialize dataset (download & filter)
|
|
assert_parity_absolute(
|
|
lightning['durations'][1:], vanilla['durations'][1:], norm_by=num_epochs, max_diff=max_diff_speed
|
|
)
|
|
|
|
assert_parity_relative(lightning['memory'], vanilla['memory'], max_diff=max_diff_memory)
|
|
|
|
|
|
def _hook_memory():
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
used_memory = torch.cuda.max_memory_allocated()
|
|
else:
|
|
used_memory = np.nan
|
|
return used_memory
|
|
|
|
|
|
def measure_loops(cls_model, kind, num_runs=10, num_epochs=10):
|
|
"""
|
|
Returns an array with the last loss from each epoch for each run
|
|
"""
|
|
hist_losses = []
|
|
hist_durations = []
|
|
hist_memory = []
|
|
|
|
device_type = "cuda" if torch.cuda.is_available() else "cpu"
|
|
torch.backends.cudnn.deterministic = True
|
|
for i in tqdm(range(num_runs), desc=f'{kind} with {cls_model.__name__}'):
|
|
gc.collect()
|
|
if device_type == 'cuda':
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.reset_max_memory_cached()
|
|
torch.cuda.reset_max_memory_allocated()
|
|
torch.cuda.reset_accumulated_memory_stats()
|
|
torch.cuda.reset_peak_memory_stats()
|
|
time.sleep(1)
|
|
|
|
time_start = time.perf_counter()
|
|
|
|
_loop = lightning_loop if kind == "PT Lightning" else vanilla_loop
|
|
final_loss, used_memory = _loop(cls_model, idx=i, device_type=device_type, num_epochs=num_epochs)
|
|
|
|
time_end = time.perf_counter()
|
|
|
|
hist_losses.append(final_loss)
|
|
hist_durations.append(time_end - time_start)
|
|
hist_memory.append(used_memory)
|
|
|
|
return {
|
|
'losses': hist_losses,
|
|
'durations': hist_durations,
|
|
'memory': hist_memory,
|
|
}
|
|
|
|
|
|
def vanilla_loop(cls_model, idx, device_type: str = 'cuda', num_epochs=10):
|
|
device = torch.device(device_type)
|
|
# set seed
|
|
seed_everything(idx)
|
|
|
|
# init model parts
|
|
model = cls_model()
|
|
dl = model.train_dataloader()
|
|
optimizer = model.configure_optimizers()
|
|
|
|
# model to GPU
|
|
model = model.to(device)
|
|
|
|
epoch_losses = []
|
|
# as the first run is skipped, no need to run it long
|
|
for epoch in range(num_epochs if idx > 0 else 1):
|
|
|
|
# run through full training set
|
|
for j, batch in enumerate(dl):
|
|
batch = [x.to(device) for x in batch]
|
|
loss_dict = model.training_step(batch, j)
|
|
loss = loss_dict['loss']
|
|
loss.backward()
|
|
optimizer.step()
|
|
optimizer.zero_grad()
|
|
|
|
# track last epoch loss
|
|
epoch_losses.append(loss.item())
|
|
|
|
return epoch_losses[-1], _hook_memory()
|
|
|
|
|
|
def lightning_loop(cls_model, idx, device_type: str = 'cuda', num_epochs=10):
|
|
seed_everything(idx)
|
|
|
|
model = cls_model()
|
|
# init model parts
|
|
trainer = Trainer(
|
|
# as the first run is skipped, no need to run it long
|
|
max_epochs=num_epochs if idx > 0 else 1,
|
|
progress_bar_refresh_rate=0,
|
|
weights_summary=None,
|
|
gpus=1 if device_type == 'cuda' else 0,
|
|
checkpoint_callback=False,
|
|
deterministic=True,
|
|
logger=False,
|
|
replace_sampler_ddp=False,
|
|
)
|
|
trainer.fit(model)
|
|
|
|
return trainer.train_loop.running_loss.last().item(), _hook_memory()
|