mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-08 09:32:20 +00:00
64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang)
|
|
#
|
|
# See ../../LICENSE for clarification regarding multiple authors
|
|
#
|
|
# 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 pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from icefall.checkpoint import average_checkpoints, load_checkpoint, save_checkpoint
|
|
|
|
|
|
@pytest.fixture
|
|
def checkpoints1(tmp_path):
|
|
f = tmp_path / "f.pt"
|
|
m = nn.Module()
|
|
m.p1 = nn.Parameter(torch.tensor([10.0, 20.0]), requires_grad=False)
|
|
m.register_buffer("p2", torch.tensor([10, 100]))
|
|
|
|
params = {"a": 10, "b": 20}
|
|
save_checkpoint(f, m, params=params)
|
|
return f
|
|
|
|
|
|
@pytest.fixture
|
|
def checkpoints2(tmp_path):
|
|
f = tmp_path / "f2.pt"
|
|
m = nn.Module()
|
|
m.p1 = nn.Parameter(torch.Tensor([50, 30.0]))
|
|
m.register_buffer("p2", torch.tensor([1, 3]))
|
|
params = {"a": 100, "b": 200}
|
|
|
|
save_checkpoint(f, m, params=params)
|
|
return f
|
|
|
|
|
|
def test_load_checkpoints(checkpoints1):
|
|
m = nn.Module()
|
|
m.p1 = nn.Parameter(torch.Tensor([0, 0.0]))
|
|
m.p2 = nn.Parameter(torch.Tensor([0, 0]))
|
|
params = load_checkpoint(checkpoints1, m)
|
|
assert torch.allclose(m.p1, torch.Tensor([10.0, 20]))
|
|
assert params["a"] == 10
|
|
assert params["b"] == 20
|
|
|
|
|
|
def test_average_checkpoints(checkpoints1, checkpoints2):
|
|
state_dict = average_checkpoints([checkpoints1, checkpoints2])
|
|
assert torch.allclose(state_dict["p1"], torch.Tensor([30, 25.0]))
|
|
assert torch.allclose(state_dict["p2"], torch.tensor([5, 51]))
|