Hyper

!apt-get install swig
!pip install -U pymvpa2
import mvpa2.datasets
import mvpa2.algorithms.hyperalignment
import numpy as np
import pickle

def hyperalignment(a, b):    
    x = [a, b]    
    dataset = [mvpa2.datasets.Dataset(x_) for x_ in x]    
    hyperalignment = mvpa2.algorithms.hyperalignment.Hyperalignment()(dataset)    
    y = [hyperalignment[j].forward(dataset[j]).samples for j in range(len(dataset))]    
    return (y[0] + y[1]) / 2

with open("yourpath/data_1.dat", 'rb') as f:
    X_tr1, _, X_te1, _ = pickle.load(f)
with open("yourpath/data_2.dat", 'rb') as f:
    X_tr2, _, X_te2, _ = pickle.load(f)
X_1 = np.array(X_te1 + X_tr1)
X_2 = np.array(X_te2 + X_tr2)
X_hyperaligned = hyperalignment(X_1, X_2)

!pip install mxnet-cu101
from __future__ import annotations
import os
from typing import Tuple, Union
import matplotlib.pyplot as plt
import mxnet as mx
from mxnet import autograd, gluon, nd, symbol
from mxnet.gluon.nn import Conv2D, Dense, HybridBlock,       
    HybridSequential, LeakyReLU
from mxnet.gluon.parameter import Parameter
from mxnet.initializer import Zero
from mxnet.io import NDArrayIter
from PIL import Image
from scipy.stats import zscore

class Pixelnorm(HybridBlock):
    def __init__(self, epsilon: float = 1e-8) -> None:
        super(Pixelnorm, self).__init__()
        self.eps = epsilon
    def hybrid_forward(self, F, x) -> nd:
        return x * F.rsqrt(F.mean(F.square(x), 1, True) + self.eps)

class Bias(HybridBlock):
    def __init__(self, shape: Tuple) -> None:
        super(Bias, self).__init__()
        self.shape = shape
        with self.name_scope():
            self.b = self.params.get("b", init=Zero(), shape=shape)
    def hybrid_forward(self, F, x, b) -> nd:
        return F.broadcast_add(x, b[None, :, None, None])

class Block(HybridSequential):
    def __init__(self, channels: int, in_channels: int) -> None:
        super(Block, self).__init__()
        self.channels = channels
        self.in_channels = in_channels
        with self.name_scope():
            self.add(Conv2D(channels, 3, padding=1, in_channels=in_channels))
            self.add(LeakyReLU(0.2))
            self.add(Pixelnorm())
            self.add(Conv2D(channels, 3, padding=1, in_channels=channels))
            self.add(LeakyReLU(0.2))
            self.add(Pixelnorm())
    def hybrid_forward(self, F, x) -> nd:
        x = F.repeat(x, 2, 2)
        x = F.repeat(x, 2, 3)
        for i in range(len(self)):
            x = self[i](x)
        return x

class Generator(HybridSequential):
    def __init__(self) -> None:
        super(Generator, self).__init__()
        with self.name_scope():
            self.add(Pixelnorm())
            self.add(Dense(8192, use_bias=False, in_units=512))
            self.add(Bias((512,)))
            self.add(LeakyReLU(0.2))
            self.add(Pixelnorm())
            self.add(Conv2D(512, 3, padding=1, in_channels=512))
            self.add(LeakyReLU(0.2))
            self.add(Pixelnorm())
            self.add(Block(512, 512))
            self.add(Block(512, 512))
            self.add(Block(512, 512))
            self.add(Block(256, 512))
            self.add(Block(128, 256))
            self.add(Block(64, 128))
            self.add(Block(32, 64))
            self.add(Block(16, 32))
            self.add(Conv2D(3, 1, in_channels=16))
    def hybrid_forward(self, F: Union(nd, symbol), x: nd, layer: int) -> nd:
        x = F.Reshape(self[1](self[0](x)), (-1, 512, 4, 4))
        for i in range(2, len(self)):
            x = self[i](x)
            if i == layer + 7:
                return x
        return x

class Linear(HybridSequential):
    def __init__(self, n_in, n_out):
        super(Linear, self).__init__()
        with self.name_scope():
            self.add(Dense(n_out, in_units=n_in))

# Set parameters.
batch_size = 30
max_epoch = 1500
n_lat = 512
n_vox = 4096

# Make dataset to take batches from during training.
def load_dataset(t, x, batch_size):
    return NDArrayIter({ "x": nd.stack(*x, axis=0) }, { "t": nd.stack(*t, axis=0) }, batch_size, True)

# Latents.
with open("yourpath/data_1.dat", 'rb') as f:
    _, T_tr, _, T_te = pickle.load(f)

# Z-score the brain data.
X_te = zscore(X_hyperaligned[:36])
X_tr = zscore(X_hyperaligned[36:])
train = load_dataset(nd.array(T_tr), nd.array(X_tr), batch_size)
test =  load_dataset(nd.array(T_te), nd.array(X_te), batch_size=36)

# Initialize generator.
generator = Generator()
generator.load_parameters("yourpath/generator.params")
mean_squared_error = gluon.loss.L2Loss()

# Initialize linear model.
vox_to_lat = Linear(n_vox, n_lat)
vox_to_lat.initialize()
trainer = gluon.Trainer(vox_to_lat.collect_params(), "Adam", {"learning_rate": 0.00001, "wd": 0.01})

# Training.
epoch = 0
results_tr = []
results_te = []
while epoch < max_epoch:
    train.reset()
    test.reset()
    loss_tr = 0
    loss_te = 0
    count = 0
    for batch_tr in train:
        with autograd.record():
            lat_Y = vox_to_lat(batch_tr.data[0])
            loss = mean_squared_error(lat_Y, batch_tr.label[0])
        loss.backward()
        trainer.step(batch_size)
        loss_tr += loss.mean().asnumpy()
        count += 1
    for batch_te in test:
        lat_Y = vox_to_lat(batch_te.data[0])
        loss = mean_squared_error(lat_Y, batch_te.label[0])
        loss_te += loss.mean().asnumpy()
    loss_tr_normalized = loss_tr / count
    results_tr.append(loss_tr_normalized)
    results_te.append(loss_te)
    epoch += 1
    print("Epoch %i: %.4f / %.4f" % (epoch, loss_tr_normalized, loss_te))
plt.figure()
plt.plot(np.linspace(0, epoch, epoch), results_tr)
plt.plot(np.linspace(0, epoch, epoch), results_te)
plt.show()

# Testing and reconstructing
lat_Y = vox_to_lat(nd.array(X_te))
dir = "yourpath/reconstructions"
if not os.path.exists(dir):
    os.mkdir(dir)
for i, latent in enumerate(lat_Y):
    face = generator(latent[None], 9).asnumpy()
    face = np.clip(np.rint(127.5 * face + 127.5), 0.0, 255.0)
    face = face.astype("uint8").transpose(0, 2, 3, 1)
    Image.fromarray(face[0], 'RGB').save(dir + "/%d.png" % i)