101 lines
3.8 KiB
Python
101 lines
3.8 KiB
Python
import pathlib
|
|
import numpy as np
|
|
|
|
from tinygrad.tensor import Tensor
|
|
|
|
def save_scalar(s, name, path):
|
|
s = np.array([1.0, float(s)]).astype(np.float32)
|
|
np.save(pathlib.Path(path, f'{name}.npy'), s)
|
|
|
|
def save_tensor(tensor, name, path):
|
|
tensor_numpy = tensor.numpy()
|
|
tensor_dims = np.array(tensor_numpy.shape)
|
|
tensor_values = tensor_numpy.flatten()
|
|
tensor_to_save = np.concatenate((tensor_dims, tensor_values)).astype(np.float32)
|
|
np.save(pathlib.Path(path, f'{name}.npy'), tensor_to_save)
|
|
|
|
def save_linear(linear, path):
|
|
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
|
save_tensor(linear.weight.transpose(), 'weight', path) # PyTorch and Tinygrad strangely transpose linear weights so reverse that
|
|
if linear.bias is not None:
|
|
save_tensor(linear.bias, 'bias', path)
|
|
|
|
def save_layer_norm(layer_norm, path):
|
|
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
|
save_tensor(layer_norm.weight, 'weight', path)
|
|
save_tensor(layer_norm.bias, 'bias', path)
|
|
save_scalar(layer_norm.eps, 'eps', path)
|
|
|
|
def save_group_norm(layer_norm, path):
|
|
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
|
if layer_norm.weight is not None:
|
|
save_tensor(layer_norm.weight, 'weight', path)
|
|
if layer_norm.bias is not None:
|
|
save_tensor(layer_norm.bias, 'bias', path)
|
|
save_scalar(layer_norm.eps, 'eps', path)
|
|
save_scalar(layer_norm.num_groups, 'n_group', path)
|
|
save_scalar(layer_norm.num_channels, 'n_channel', path)
|
|
|
|
def to_tuple_tensor(val):
|
|
if isinstance(val, tuple):
|
|
# Convert tuple to Tensor
|
|
if len(val) == 1:
|
|
return Tensor([val[0], val[0]])
|
|
elif len(val) == 2:
|
|
return Tensor([val[0], val[1]])
|
|
else:
|
|
raise ValueError('Tuple should be of length 1 or 2 only.')
|
|
else:
|
|
# Treat as scalar and convert to Tensor
|
|
return Tensor([val, val])
|
|
|
|
def save_conv2d(conv2d, path):
|
|
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
|
|
|
save_tensor(conv2d.weight, 'weight', path)
|
|
if conv2d.bias is not None:
|
|
save_tensor(conv2d.bias, 'bias', path)
|
|
save_tensor(to_tuple_tensor(conv2d.stride), 'stride', path)
|
|
save_tensor(to_tuple_tensor(conv2d.padding), 'padding', path)
|
|
save_tensor(to_tuple_tensor(conv2d.dilation), 'dilation', path)
|
|
save_scalar(conv2d.groups, "n_group", path)
|
|
save_tensor(to_tuple_tensor(conv2d.kernel_size), 'kernel_size', path)
|
|
|
|
assert conv2d.groups == 1
|
|
in_channels = conv2d.weight.shape[1]
|
|
out_channels = conv2d.weight.shape[0]
|
|
save_scalar(in_channels, "n_channels_in", path)
|
|
save_scalar(out_channels, "n_channels_out", path)
|
|
|
|
def save_padded_conv2d(padded_conv2d, path):
|
|
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
|
|
|
# Store conv2d layer weights
|
|
orig_padding = padded_conv2d.padding
|
|
padded_conv2d.padding = (0, 0)
|
|
save_conv2d(padded_conv2d, f"{path}/conv")
|
|
padded_conv2d.padding = orig_padding
|
|
|
|
# Dimensions: in-channels and out-channels
|
|
assert padded_conv2d.groups == 1
|
|
channels = (padded_conv2d.weight.shape[1], padded_conv2d.weight.shape[0])
|
|
save_tensor(to_tuple_tensor(channels), 'channels', path)
|
|
|
|
assert len(padded_conv2d.kernel_size) == 1 or padded_conv2d.kernel_size[0] == padded_conv2d.kernel_size[1]
|
|
save_scalar(padded_conv2d.kernel_size[0], 'kernel_size', path)
|
|
|
|
# Stride
|
|
assert not isinstance(padded_conv2d.stride, tuple) or len(padded_conv2d.stride) == 1
|
|
save_scalar(padded_conv2d.stride, 'stride', path)
|
|
|
|
# Padding
|
|
padding = [padded_conv2d.padding[0], padded_conv2d.padding[1],
|
|
padded_conv2d.padding[2], padded_conv2d.padding[3]]
|
|
save_tensor(Tensor(padding), 'padding', path)
|
|
|
|
|
|
def save_embedding(embedding, path):
|
|
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
|
save_tensor(embedding.weight, 'weight', path)
|
|
|