mmcv/tests/test_runner.py
Kai Chen eae81c1e86
Add yapf and isort to travis (#96)
* add yapf and isort to travis

* minor formatting

* remove blank lines

* skip unit tests for progressbar when python2

* update travis to ubuntu 16.04

* use a newer version ffmpeg

* add -y to add-apt-repository
2019-07-30 23:15:56 +08:00

30 lines
747 B
Python

import os.path as osp
import tempfile
import warnings
def test_save_checkpoint():
try:
import torch
import torch.nn as nn
except ImportError:
warnings.warn('Skipping test_save_checkpoint in the absense of torch')
return
import mmcv.runner
model = nn.Linear(1, 1)
runner = mmcv.runner.Runner(model=model, batch_processor=lambda x: x)
with tempfile.TemporaryDirectory() as root:
runner.save_checkpoint(root)
latest_path = osp.join(root, 'latest.pth')
epoch1_path = osp.join(root, 'epoch_1.pth')
assert osp.exists(latest_path)
assert osp.exists(epoch1_path)
assert osp.realpath(latest_path) == epoch1_path
torch.load(latest_path)