<reponame>zhengzangw/Fed-SINGA
import json
import os
import time
import zipfile

import numpy as np
from singa import device, opt, tensor

from .data import mnist
from .model import cnn
from .train import resize_dataset

# save load states constant
TENSOR_DICT_FILENAME = "/tensor_dict.npz"
STATES_ATTR_FILENAME = "/states_attr.json"
MODEL_STATE_TYPE = 0
AUX_STATE_TYPE = 1
CENTRAL_MODEL_PATH = "checkpoint/central_model.zip"

np_dtype = {"float16": np.float16, "float32": np.float32}

singa_dtype = {"float16": tensor.float16, "float32": tensor.float32}


# Calculate accuracy
def accuracy(pred, target):
    # y is network output to be compared with ground truth (int)
    y = np.argmax(pred, axis=1)
    a = y == target
    correct = np.array(a, "int").sum()
    return correct


def load_states(fpath):
    assert os.path.isfile(fpath), "Failed to load states, %s is not exist." % fpath

    timestamp = time.time()
    tmp_dir = "/tmp/singa_load_states_%s" % timestamp
    os.mkdir(tmp_dir)

    with zipfile.ZipFile(fpath, "r") as zf:
        zf.extractall(tmp_dir)

    tensor_dict_fp = tmp_dir + TENSOR_DICT_FILENAME
    states_attr_fp = tmp_dir + STATES_ATTR_FILENAME

    with open(states_attr_fp) as f:
        states_attr = json.load(f)

    tensor_dict = np.load(tensor_dict_fp)

    # restore singa tensor from numpy
    model_states = dict()
    aux_states = dict()

    for k in tensor_dict.files:
        if states_attr[k]["state_type"] == MODEL_STATE_TYPE:
            model_states[k] = tensor.from_numpy(tensor_dict[k])
        elif states_attr[k]["state_type"] == AUX_STATE_TYPE:
            aux_states[k] = tensor.from_numpy(tensor_dict[k])

    # clean up tmp files
    os.remove(tensor_dict_fp)
    os.remove(states_attr_fp)
    os.rmdir(tmp_dir)
    return model_states, aux_states


def save_states(fpath, states, aux_states={}):
    """Save states.
    Args:
        fpath: output file path (without the extension)
        aux_states(dict): values are standard data types or Tensor,
                          e.g., epoch ID, learning rate, optimizer states
    """
    assert not os.path.isfile(fpath), "Failed to save states, %s is already existed." % fpath

    # save states data and attr
    tensor_dict = {}
    states_attr = {}
    for k, v in states.items():
        assert isinstance(v, tensor.Tensor), "Only tensor state is allowed"
        tensor_dict[k] = tensor.to_numpy(v)
        states_attr[k] = {"state_type": MODEL_STATE_TYPE, "shape": v.shape, "dtype": v.dtype}

    for k, v in aux_states.items():
        assert isinstance(v, tensor.Tensor), "Only tensor aux state is allowed"
        tensor_dict[k] = tensor.to_numpy(v)
        states_attr[k] = {"state_type": AUX_STATE_TYPE, "shape": v.shape, "dtype": v.dtype}

    # save to files
    timestamp = time.time()
    tmp_dir = "/tmp/singa_save_states_%s" % timestamp
    os.mkdir(tmp_dir)
    tensor_dict_fp = tmp_dir + TENSOR_DICT_FILENAME
    states_attr_fp = tmp_dir + STATES_ATTR_FILENAME

    np.savez(tensor_dict_fp, **tensor_dict)

    with open(states_attr_fp, "w") as fp:
        json.dump(states_attr, fp)

    compression = zipfile.ZIP_DEFLATED
    with zipfile.ZipFile(fpath, mode="w") as zf:
        zf.write(tensor_dict_fp, os.path.basename(tensor_dict_fp), compress_type=compression)
        zf.write(states_attr_fp, os.path.basename(states_attr_fp), compress_type=compression)

    # clean up tmp files
    os.remove(tensor_dict_fp)
    os.remove(states_attr_fp)
    os.rmdir(tmp_dir)


def eval():
    dev = device.get_default_device()
    batch_size = 64
    precision = "float32"
    world_size = 1

    train_x, train_y, val_x, val_y = mnist.load()
    num_channels = train_x.shape[1]
    image_size = train_x.shape[2]
    data_size = np.prod(train_x.shape[1 : train_x.ndim]).item()
    num_classes = (np.max(train_y) + 1).item()

    model = cnn.create_model(num_channels=num_channels, num_classes=num_classes)

    if model.dimension == 4:
        tx = tensor.Tensor(
            (batch_size, num_channels, model.input_size, model.input_size),
            dev,
            singa_dtype[precision],
        )
    elif model.dimension == 2:
        tx = tensor.Tensor((batch_size, data_size), dev, singa_dtype[precision])
        np.reshape(train_x, (train_x.shape[0], -1))
        np.reshape(val_x, (val_x.shape[0], -1))

    # Attach model to graph
    sgd = opt.SGD(lr=0.005, momentum=0.9, weight_decay=1e-5, dtype=singa_dtype[precision])
    model.set_optimizer(sgd)
    model.compile([tx], is_train=True, use_graph=True, sequential=False)
    dev.SetVerbosity(False)
    # print("loading model from " + CENTRAL_MODEL_PATH)
    model.load_states(fpath=CENTRAL_MODEL_PATH)

    # print(model.get_states()['linear2.b'])

    ty = tensor.Tensor((batch_size,), dev, tensor.int32)
    num_val_batch = val_x.shape[0] // batch_size
    test_correct = np.zeros(shape=[1], dtype=np.float32)
    model.eval()

    for b in range(num_val_batch):
        x = val_x[b * batch_size : (b + 1) * batch_size]
        if model.dimension == 4:
            if image_size != model.input_size:
                x = resize_dataset(x, model.input_size)
        x = x.astype(np_dtype[precision])
        y = val_y[b * batch_size : (b + 1) * batch_size]
        tx.copy_from_numpy(x)
        ty.copy_from_numpy(y)
        out_test = model(tx)
        test_correct += accuracy(tensor.to_numpy(out_test), y)

    # Output the evaluation accuracy
    print("Evaluation accuracy = %f" % (test_correct / (num_val_batch * batch_size * world_size)))


def aggregate(num):
    model_states = {}
    for d in range(num):
        checkpointpath = "checkpoint/checkpoint_" + str(d) + ".zip"
        tmp, aux_states = load_states(checkpointpath)
        if not bool(model_states):
            model_states = tmp
        else:
            for k, v in tmp.items():
                model_states[k] = model_states[k] + v

    for k, v in model_states.items():
        model_states[k] = model_states[k] / num

    if os.path.exists(CENTRAL_MODEL_PATH):
        os.remove(CENTRAL_MODEL_PATH)
    save_states(CENTRAL_MODEL_PATH, model_states)


aggregate(10)
eval()

import json
import os
import time
import zipfile

import numpy as np
from singa import device, opt, tensor

from .data import mnist
from .model import cnn
from .train import resize_dataset

# save load states constant
TENSOR_DICT_FILENAME = "/tensor_dict.npz"
STATES_ATTR_FILENAME = "/states_attr.json"
MODEL_STATE_TYPE = 0
AUX_STATE_TYPE = 1
CENTRAL_MODEL_PATH = "checkpoint/central_model.zip"

np_dtype = {"float16": np.float16, "float32": np.float32}

singa_dtype = {"float16": tensor.float16, "float32": tensor.float32}


# Calculate accuracy
def accuracy(pred, target):
    # y is network output to be compared with ground truth (int)
    y = np.argmax(pred, axis=1)
    a = y == target
    correct = np.array(a, "int").sum()
    return correct


def load_states(fpath):
    assert os.path.isfile(fpath), "Failed to load states, %s is not exist." % fpath

    timestamp = time.time()
    tmp_dir = "/tmp/singa_load_states_%s" % timestamp
    os.mkdir(tmp_dir)

    with zipfile.ZipFile(fpath, "r") as zf:
        zf.extractall(tmp_dir)

    tensor_dict_fp = tmp_dir + TENSOR_DICT_FILENAME
    states_attr_fp = tmp_dir + STATES_ATTR_FILENAME

    with open(states_attr_fp) as f:
        states_attr = json.load(f)

    tensor_dict = np.load(tensor_dict_fp)

    # restore singa tensor from numpy
    model_states = dict()
    aux_states = dict()

    for k in tensor_dict.files:
        if states_attr[k]["state_type"] == MODEL_STATE_TYPE:
            model_states[k] = tensor.from_numpy(tensor_dict[k])
        elif states_attr[k]["state_type"] == AUX_STATE_TYPE:
            aux_states[k] = tensor.from_numpy(tensor_dict[k])

    # clean up tmp files
    os.remove(tensor_dict_fp)
    os.remove(states_attr_fp)
    os.rmdir(tmp_dir)
    return model_states, aux_states


def save_states(fpath, states, aux_states={}):
    """Save states.
    Args:
        fpath: output file path (without the extension)
        aux_states(dict): values are standard data types or Tensor,
                          e.g., epoch ID, learning rate, optimizer states
    """
    assert not os.path.isfile(fpath), "Failed to save states, %s is already existed." % fpath

    # save states data and attr
    tensor_dict = {}
    states_attr = {}
    for k, v in states.items():
        assert isinstance(v, tensor.Tensor), "Only tensor state is allowed"
        tensor_dict[k] = tensor.to_numpy(v)
        states_attr[k] = {"state_type": MODEL_STATE_TYPE, "shape": v.shape, "dtype": v.dtype}

    for k, v in aux_states.items():
        assert isinstance(v, tensor.Tensor), "Only tensor aux state is allowed"
        tensor_dict[k] = tensor.to_numpy(v)
        states_attr[k] = {"state_type": AUX_STATE_TYPE, "shape": v.shape, "dtype": v.dtype}

    # save to files
    timestamp = time.time()
    tmp_dir = "/tmp/singa_save_states_%s" % timestamp
    os.mkdir(tmp_dir)
    tensor_dict_fp = tmp_dir + TENSOR_DICT_FILENAME
    states_attr_fp = tmp_dir + STATES_ATTR_FILENAME

    np.savez(tensor_dict_fp, **tensor_dict)

    with open(states_attr_fp, "w") as fp:
        json.dump(states_attr, fp)

    compression = zipfile.ZIP_DEFLATED
    with zipfile.ZipFile(fpath, mode="w") as zf:
        zf.write(tensor_dict_fp, os.path.basename(tensor_dict_fp), compress_type=compression)
        zf.write(states_attr_fp, os.path.basename(states_attr_fp), compress_type=compression)

    # clean up tmp files
    os.remove(tensor_dict_fp)
    os.remove(states_attr_fp)
    os.rmdir(tmp_dir)


def eval():
    dev = device.get_default_device()
    batch_size = 64
    precision = "float32"
    world_size = 1

    train_x, train_y, val_x, val_y = mnist.load()
    num_channels = train_x.shape[1]
    image_size = train_x.shape[2]
    data_size = np.prod(train_x.shape[1 : train_x.ndim]).item()
    num_classes = (np.max(train_y) + 1).item()

    model = cnn.create_model(num_channels=num_channels, num_classes=num_classes)

    if model.dimension == 4:
        tx = tensor.Tensor(
            (batch_size, num_channels, model.input_size, model.input_size),
            dev,
            singa_dtype[precision],
        )
    elif model.dimension == 2:
        tx = tensor.Tensor((batch_size, data_size), dev, singa_dtype[precision])
        np.reshape(train_x, (train_x.shape[0], -1))
        np.reshape(val_x, (val_x.shape[0], -1))

    # Attach model to graph
    sgd = opt.SGD(lr=0.005, momentum=0.9, weight_decay=1e-5, dtype=singa_dtype[precision])
    model.set_optimizer(sgd)
    model.compile([tx], is_train=True, use_graph=True, sequential=False)
    dev.SetVerbosity(False)
    # print("loading model from " + CENTRAL_MODEL_PATH)
    model.load_states(fpath=CENTRAL_MODEL_PATH)

    # print(model.get_states()['linear2.b'])

    ty = tensor.Tensor((batch_size,), dev, tensor.int32)
    num_val_batch = val_x.shape[0] // batch_size
    test_correct = np.zeros(shape=[1], dtype=np.float32)
    model.eval()

    for b in range(num_val_batch):
        x = val_x[b * batch_size : (b + 1) * batch_size]
        if model.dimension == 4:
            if image_size != model.input_size:
                x = resize_dataset(x, model.input_size)
        x = x.astype(np_dtype[precision])
        y = val_y[b * batch_size : (b + 1) * batch_size]
        tx.copy_from_numpy(x)
        ty.copy_from_numpy(y)
        out_test = model(tx)
        test_correct += accuracy(tensor.to_numpy(out_test), y)

    # Output the evaluation accuracy
    print("Evaluation accuracy = %f" % (test_correct / (num_val_batch * batch_size * world_size)))


def aggregate(num):
    model_states = {}
    for d in range(num):
        checkpointpath = "checkpoint/checkpoint_" + str(d) + ".zip"
        tmp, aux_states = load_states(checkpointpath)
        if not bool(model_states):
            model_states = tmp
        else:
            for k, v in tmp.items():
                model_states[k] = model_states[k] + v

    for k, v in model_states.items():
        model_states[k] = model_states[k] / num

    if os.path.exists(CENTRAL_MODEL_PATH):
        os.remove(CENTRAL_MODEL_PATH)
    save_states(CENTRAL_MODEL_PATH, model_states)


aggregate(10)
eval()
# save load states constant # Calculate accuracy # y is network output to be compared with ground truth (int) # restore singa tensor from numpy # clean up tmp files Save states.     Args:         fpath: output file path (without the extension)         aux_states(dict): values are standard data types or Tensor,                           e.g., epoch ID, learning rate, optimizer states # save states data and attr # save to files # clean up tmp files # Attach model to graph # print("loading model from " + CENTRAL_MODEL_PATH) # print(model.get_states()['linear2.b']) # Output the evaluation accuracy