File size: 2,990 Bytes
ec188ed | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | import argparse
from pathlib import Path
import torch
from model import Generator
def _read_tf_checkpoint(checkpoint_path):
try:
import tensorflow as tf
except ImportError as exc:
raise RuntimeError("TensorFlow is required only for checkpoint conversion.") from exc
return tf.train.load_checkpoint(str(checkpoint_path))
def _copy_conv(reader, state, tf_name, torch_name):
state[torch_name + ".weight"] = torch.from_numpy(reader.get_tensor(tf_name + "/kernel")).permute(3, 2, 0, 1)
state[torch_name + ".bias"] = torch.from_numpy(reader.get_tensor(tf_name + "/bias"))
def _copy_deconv(reader, state, tf_name, torch_name):
state[torch_name + ".weight"] = torch.from_numpy(reader.get_tensor(tf_name + "/kernel")).permute(3, 2, 0, 1)
state[torch_name + ".bias"] = torch.from_numpy(reader.get_tensor(tf_name + "/bias"))
def _copy_bn(reader, state, tf_name, torch_name):
state[torch_name + ".weight"] = torch.from_numpy(reader.get_tensor(tf_name + "/gamma"))
state[torch_name + ".bias"] = torch.from_numpy(reader.get_tensor(tf_name + "/beta"))
state[torch_name + ".running_mean"] = torch.from_numpy(reader.get_tensor(tf_name + "/moving_mean"))
state[torch_name + ".running_var"] = torch.from_numpy(reader.get_tensor(tf_name + "/moving_variance"))
def convert(checkpoint_path, output_path):
checkpoint_path = Path(checkpoint_path)
output_path = Path(output_path)
reader = _read_tf_checkpoint(checkpoint_path)
generator = Generator()
state = generator.state_dict()
_copy_conv(reader, state, "generator/g_conv0/conv2d", "down_convs.0")
for i in range(1, 8):
_copy_conv(reader, state, "generator/g_conv%d/conv2d" % i, "down_convs.%d" % i)
_copy_bn(
reader,
state,
"generator/g_conv%d/batch_normalization" % i,
"down_bns.%d" % (i - 1),
)
for i in range(8):
_copy_deconv(reader, state, "generator/g_deconv%d/conv2d_transpose" % i, "up_convs.%d" % i)
if i < 7:
_copy_bn(
reader,
state,
"generator/g_deconv%d/batch_normalization" % i,
"up_bns.%d" % i,
)
generator.load_state_dict(state)
output_path.parent.mkdir(parents=True, exist_ok=True)
torch.save({"state_dict": generator.state_dict()}, output_path)
return output_path
def main():
parser = argparse.ArgumentParser(description="Convert StarNet TF1 generator checkpoint to PyTorch")
parser.add_argument("--checkpoint", default="model.ckpt", help="TensorFlow checkpoint prefix")
parser.add_argument("--output", default="starnet_generator.pth", help="Output PyTorch weights")
args = parser.parse_args()
output_path = convert(args.checkpoint, args.output)
print("Saved PyTorch weights to %s" % output_path)
if __name__ == "__main__":
main()
|