def __init__(self, X, arm_sketches, ranks):
tl.set_backend('numpy')
self.arms = []
self.core_tensor = None
self.X = X
self.arm_sketches = arm_sketches
self.ranks = ranks
def fit(
self,
*views: np.ndarray,
):
self.n_views = len(views)
self.check_params()
assert (len(
self.c) == len(views)), 'c requires as many values as #views'
train_views, covs_invsqrt = self.setup_tensor(*views)
for i, el in enumerate(train_views):
if i == 0:
M = el
else:
for _ in range(len(M.shape) - 1):
el = np.expand_dims(el, 1)
M = np.expand_dims(M, -1) @ el
M = np.mean(M, 0)
tl.set_backend('numpy')
M_parafac = parafac(M, self.latent_dims, verbose=True)
self.alphas = [
cov_invsqrt @ fac for i, (view, cov_invsqrt, fac) in enumerate(
zip(train_views, covs_invsqrt, M_parafac.factors))
]
self.score_list = [
view @ self.alphas[i] for i, view in enumerate(train_views)
]
self.weights_list = [
weights / np.linalg.norm(score)
for weights, score in zip(self.alphas, self.score_list)
]
self.score_list = [
view @ self.weights_list[i] for i, view in enumerate(train_views)
]
self.train_correlations = self.predict_corr(*views)
return self
def fit(self, *views: Tuple[np.ndarray, ...], ):
if self.c is None:
self.c = [0] * len(views)
assert (len(self.c) == len(views)), 'c requires as many values as #views'
z = self.demean_data(*views)
n = z[0].shape[0]
covs = [(1 - self.c[i]) * view.T @ view / (1 - n) + self.c[i] * np.eye(view.shape[1]) for i, view in
enumerate(z)]
covs_invsqrt = [np.linalg.inv(sqrtm(cov)) for cov in covs]
z = [z_ @ cov_invsqrt for z_, cov_invsqrt in zip(z, covs_invsqrt)]
for i, el in enumerate(z):
if i == 0:
M = el
else:
for _ in range(len(M.shape) - 1):
el = np.expand_dims(el, 1)
M = np.expand_dims(M, -1) @ el
M = np.mean(M, 0)
# for i, cov_invsqrt in enumerate(covs_invsqrt):
# M = np.tensordot(M, cov_invsqrt, axes=[[0], [0]])
tl.set_backend('numpy')
M_parafac = parafac(M, self.latent_dims, verbose=True)
self.weights_list = [cov_invsqrt @ fac for i, (view, cov_invsqrt, fac) in
enumerate(zip(z, covs_invsqrt, M_parafac.factors))]
self.score_list = [view @ self.weights_list[i] for i, view in enumerate(z)]
self.weights_list = [weights / np.linalg.norm(score) for weights, score in
zip(self.weights_list, self.score_list)]
self.score_list = [view @ self.weights_list[i] for i, view in enumerate(z)]
self.train_correlations = self.predict_corr(*views)
return self
def loss(self, *views):
m = views[0].size(0)
views = _demean(*views)
covs = [
(1 - self.r) * view.T @ view
+ self.r * torch.eye(view.size(1), device=view.device)
for view in views
]
whitened_z = [
view @ mat_pow(cov, -0.5, self.eps) for view, cov in zip(views, covs)
]
# The idea here is to form a matrix with M dimensions one for each view where at index
# M[p_i,p_j,p_k...] we have the sum over n samples of the product of the pth feature of the
# ith, jth, kth view etc.
for i, el in enumerate(whitened_z):
# To achieve this we start with the first view so M is nxp.
if i == 0:
M = el
# For the remaining views we expand their dimensions to match M i.e. nx1x...x1xp
else:
for _ in range(len(M.size()) - 1):
el = torch.unsqueeze(el, 1)
# Then we perform an outer product by expanding the dimensionality of M and
# outer product with the expanded el
M = torch.unsqueeze(M, -1) @ el
M = torch.mean(M, 0)
tl.set_backend("pytorch")
M_parafac = parafac(
M.detach(), self.latent_dims, verbose=False, normalize_factors=True
)
M_parafac.weights = 1
M_hat = cp_to_tensor(M_parafac)
return torch.linalg.norm(M - M_hat)
def Tucker_for_linear(model, decompose_rate):
tl.set_backend('pytorch')
linear_flag = 1
for i, key in enumerate(model.classifier._modules.keys()):
if isinstance(model.classifier._modules[key],
torch.nn.modules.linear.Linear):
if linear_flag != 0:
print('first linear layer')
linear_layer = model.classifier._modules[key]
decomposed = tucker_for_first_linear_layer(
linear_layer, decompose_rate)
model.classifier._modules[key] = decomposed
print('linear layer decompose end\n')
linear_flag = 0
else:
linear_layer = model.classifier._modules[key]
decomposed = tucker_decomposition_linear_layer(
linear_layer, decompose_rate)
model.classifier._modules[key] = decomposed
print('linear layer decompose end\n')
torch.save(model, 'decomposed_model.pkl')
print('tucker decompose end,model saved!\n')
print(model)
return model
def Tucker_decompose(model, ranks):
tl.set_backend('pytorch')
print('tucker decompose start\n')
N = len(model.features._modules.keys())
print('N is :', N)
conv_flag = 1
conv_count = 2
for i, key in enumerate(model.features._modules.keys()):
print('present layer is :', model.features._modules[key])
if isinstance(model.features._modules[key],
torch.nn.modules.conv.Conv2d):
conv_layer = model.features._modules[key]
if conv_flag != 0:
decomposed = tucker_for_first_conv_layer(conv_layer, ranks)
model.features._modules[key] = decomposed
print('conv layer decompose end')
conv_flag = 0
else:
decomposed = tucker_decomposition_conv_layer(
conv_layer, ranks, conv_count)
model.features._modules[key] = decomposed
print('conv layer decompose end')
conv_count += 2
torch.save(model, 'decomposed_model.pkl')
print('tucker decompose end,model saved!')
# print(model)
return model
def test_set_backend_local_threadsafe():
pytest.importorskip('torch')
global_default = tl.get_backend()
with ThreadPoolExecutor(max_workers=1) as executor:
with tl.backend_context('numpy', local_threadsafe=True):
assert tl.get_backend() == 'numpy'
# Changes only happen locally in this thread
assert executor.submit(tl.get_backend).result() == global_default
# Set the global default backend
try:
tl.set_backend('pytorch', local_threadsafe=False)
# Changed toplevel default in all threads
assert executor.submit(tl.get_backend).result() == 'pytorch'
with tl.backend_context('numpy', local_threadsafe=True):
assert tl.get_backend() == 'numpy'
def check():
assert tl.get_backend() == 'pytorch'
with tl.backend_context('numpy', local_threadsafe=True):
assert tl.get_backend() == 'numpy'
assert tl.get_backend() == 'pytorch'
executor.submit(check).result()
finally:
tl.set_backend(global_default, local_threadsafe=False)
executor.submit(tl.set_backend, global_default).result()
assert tl.get_backend() == global_default
assert executor.submit(tl.get_backend).result() == global_default
def get_decomposed_tensor(i, obs_dir, ranks):
tf.enable_eager_execution()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
X = load_observation(obs_dir)
if type(X) == np.ndarray and X.shape[0] == 65:
X = np.reshape(X, (65, 100, 116, 116))
X_tf = tf.convert_to_tensor(X, dtype=tf.float32)
tl.set_backend('tensorflow')
core, factors = tucker(X_tf,
ranks=ranks,
init='random',
tol=10e-5,
verbose=False)
transition_features = core.numpy().flatten()
line = transition_features
out = line
print('Completed line {}'.format(i))
else:
print('Skipped line {}'.format(i))
out = []
sess.close()
return out
def create_cp(dims,
rank,
sparsity=None,
method='rand',
weights=False,
return_tensor=False,
noise=None,
sparse_noise=True):
# TODO: investigate performance impact of setting backend here
tl.set_backend('pytorch')
if method == 'rand':
randfunc = torch.rand
elif method == 'randn':
randfunc = torch.randn
else:
raise NotImplementedError(f'Unknown random method: {method}')
n_dims = len(dims)
factors = [randfunc((dim, rank)) for dim in dims]
if sparsity is not None:
if isinstance(sparsity, float):
sparsity = [sparsity for _ in range(n_dims)]
elif not isinstance(sparsity, list) and not isinstance(
sparsity, tuple):
raise ValueError(
'Sparsity parameter should either be a float or tuple/list.')
# Sparsify factors
for dim in range(n_dims):
n_el = dims[dim] * rank
to_del = round(sparsity[dim] * n_el)
if to_del == 0:
continue
idxs = torch.tensor(random.sample(range(n_el), to_del))
factors[dim].view(-1)[idxs] = 0
# torch.randperm(n_el, device=device)[:n_select]
ten = None
# Add noise
if noise is not None:
ten = tl.cp_to_tensor((torch.ones(rank), factors))
if (sparsity is None or not sparse_noise):
nten = torch.randn(ten.size())
ten += noise * (norm(ten) / norm(nten)) * nten
else:
flat = ten.view(-1)
nzs = torch.nonzero(flat, as_tuple=True)[0]
nvec = torch.randn(nzs.size(0))
flat[nzs] += noise * (norm(ten) / norm(nvec)) * nvec
if return_tensor:
if ten is None:
return tl.cp_to_tensor((torch.ones(rank), factors))
return ten
if weights:
return torch.ones(rank), factors
return factors
def __init__(self, shape, n_components):
# Set backend as numpy
tl.set_backend('numpy')
self.shape = shape
self.n_components = n_components
self.Z = []
self.X = None
self.rand_init()
def __init__(self, X, ranks, ks=[], ss=[], random_seed=1, store_phis=True):
tl.set_backend('numpy')
self.X = X
self.ranks = ranks
self.ks = ks
self.ss = ss
self.random_seed = random_seed
self.store_phis = store_phis
def __init__(self, sketchs, core_sketch, Tinfo_bucket, Rinfo_bucket):
tl.set_backend('numpy')
self.arms = []
self.core_tensor = None
self.sketchs = sketchs
self.tensor_shape, self.k, self.rank, self.s = Tinfo_bucket.get_info()
self.Rinfo_bucket = Rinfo_bucket
self.core_sketch = core_sketch
def __init__(self, n, rank, k, s, dim, Rinfo_bucket, gen_typ, noise_level):
tl.set_backend('numpy')
self.n, self.rank, self.k, self.s, self.dim = n, rank, k, s, dim
self.std, self.typ, self.random_seed, self.sparse_factor = Rinfo_bucket.get_info(
)
self.total_num = np.prod(np.repeat(n, dim))
self.gen_typ = gen_typ
self.noise_level = noise_level
self.Rinfo_bucket = Rinfo_bucket
def test(dataset, cwfa, action_encoder, obs_encoder):
encoded_action = action_encoder(dataset.action.reshape(-1, dataset.action.shape[-1])).\
reshape(dataset.action.shape[0], dataset.action.shape[1], -1).cpu().detach().numpy()
encoded_obs = obs_encoder(dataset.obs.reshape(-1, dataset.obs.shape[-1])).\
reshape(dataset.obs.shape[0], dataset.obs.shape[1], -1).cpu().detach().numpy()
tl.set_backend('numpy')
pred = []
for i in range(len(encoded_obs)):
pred.append(cwfa.predict(encoded_action[i], encoded_obs[i]))
return np.asarray(pred).squeeze()
def __init__(self, Tinfo_bucket, Rinfo_bucket, gen_typ, noise_level):
tl.set_backend('numpy')
self.tensor_shape, self.k, self.rank, self.s = Tinfo_bucket.get_info()
self.n = self.tensor_shape[0]
self.dim = len(self.tensor_shape)
self.std, self.typ, self.random_seed, self.sparse_factor = Rinfo_bucket.get_info(
)
self.total_num = np.prod(self.tensor_shape)
self.gen_typ = gen_typ
self.noise_level = noise_level
self.Tinfo_bucket = Tinfo_bucket
self.Rinfo_bucket = Rinfo_bucket
def main():
start_time = time.time()
args = parse_arguments()
lines, verb2id, subject2id, object2id = get_dict_and_samples(
args.input_path, args.min_count, args.first_n, args.step)
if args.sparse:
large_tensor = create_sparse_tensor(lines, verb2id, subject2id,
object2id)
print("Decomposition started")
if args.algorithm == 'tucker':
weights, factors = partial_tucker(large_tensor,
modes=[0, 1, 2],
rank=args.embedding_size,
init='random')
else:
weights, factors = sparse_parafac(large_tensor,
rank=args.embedding_size,
init='random')
else:
tl.set_backend('pytorch')
large_tensor = create_tensor(lines, verb2id, subject2id, object2id)
print("Decomposition started")
if args.algorithm == 'tucker':
weights, factors = tucker(large_tensor,
rank=args.embedding_size,
init='random')
else:
weights, factors = parafac(large_tensor,
rank=args.embedding_size,
init='random')
factors = [factor.cpu().numpy().astype(float) for factor in factors]
assert [factor.shape[0] for factor in factors
] == [len(verb2id), len(subject2id),
len(object2id)]
output_path = os.path.join(
args.output_path,
f"{args.input_path[5:13]}-{args.algorithm}_e{args.embedding_size}_"
f"min-count-{args.min_count}_cut-{args.first_n}_step-{args.step}")
if not os.path.exists(output_path):
os.mkdir(output_path)
save_to_file(factors[0], verb2id, os.path.join(output_path, 'verbs.tsv'))
save_to_file(factors[1], subject2id,
os.path.join(output_path, 'subjects.tsv'))
save_to_file(factors[2], object2id, os.path.join(output_path,
'objects.tsv'))
print(f"---- Took {(time.time() - start_time)} seconds ----")
def __init__(self, X, ks, random_seed, ss = [], typ = 'g', \
sparse_factor = 0.1, std = 1, store_phis = True):
'''
:param X: tensor being skeched
:param ks: k, the reduced dimension of the arm tensors, an 1-d array
:param ss: At any index, the element of ss is greater than the element
of kswhen ss = [], do not perform core sketch, that is, core_sketch == X
:param random_seed: random_seed
:param sparse_factor: only typ == 'sp', p matters representing the
sparse factor
'''
tl.set_backend('numpy')
self.X = X
self.N = len(X.shape)
self.ss = ss
self.ks = ks
self.typ = typ
self.sparse_factor = sparse_factor
self.arm_sketches = []
self.random_seed = random_seed
self.core_sketch = X
self.tensor_shape = X.shape
self.phis = []
self.std = std
# set the random seed for following procedure
np.random.seed(random_seed)
Rinfo_bucket = RandomInfoBucket(std = self.std, typ=self.typ,
random_seed = self.random_seed, sparse_factor = self.sparse_factor)
rm_generator = Sketch.sketch_arm_rm_generator(self.tensor_shape, \
self.ks, Rinfo_bucket)
mode_n = 0
for rm in rm_generator:
self.arm_sketches.append(np.dot(tl.unfold(self.X, mode=mode_n), rm))
mode_n += 1
np.random.seed(random_seed)
if self.ss != []:
rm_generator = Sketch.sketch_core_rm_generator(self.tensor_shape, \
self.ss, Rinfo_bucket)
mode_n = 0
for rm in rm_generator:
self.phis.append(rm)
self.core_sketch = tl.tenalg.mode_dot(self.core_sketch, rm,\
mode=mode_n)
mode_n += 1
if not store_phis:
self.phis = []
def test_cpd_als_tensorly(benchmark):
for datatype in BACKEND_TYPES:
tl.set_backend(datatype)
assert tl.get_backend() == datatype
_, input_tensor_val = init_rand_cp(dim, size, rank)
input_tensor = tl.tensor(input_tensor_val, dtype='float64')
factors = benchmark(parafac,
input_tensor,
rank=rank,
init='random',
tol=0,
n_iter_max=1,
verbose=0)
def __init__(self,
arm_sketches,
core_sketch,
Tinfo_bucket,
Rinfo_bucket,
phis=[]):
tl.set_backend('numpy')
self.arms = []
self.core_tensor = None
self.arm_sketches = arm_sketches
# Note get_info extract some extraneous information
self.tensor_shape, self.ks, self.ranks, self.ss = Tinfo_bucket.get_info(
)
self.Rinfo_bucket = Rinfo_bucket
self.phis = phis
self.core_sketch = core_sketch
def prox_nuclear_truncated_2(self, data, alpha, k=50):
import tensorly as tl
tl.set_backend('pytorch')
U, S, V = tl.truncated_svd(data.cpu(), n_eigenvecs=k)
U, S, V = torch.FloatTensor(U).cuda(), torch.FloatTensor(S).cuda(), torch.FloatTensor(V).cuda()
self.nuclear_norm = S.sum()
# print("nuclear norm: %.4f" % self.nuclear_norm)
S = torch.clamp(S-alpha, min=0)
indices = torch.tensor(range(0, U.shape[0]),range(0, U.shape[0])).cuda()
values = S
diag_S = torch.sparse.FloatTensor(indices, values, torch.Size(U.shape))
# diag_S = torch.diag(torch.clamp(S-alpha, min=0))
U = torch.spmm(U, diag_S)
V = torch.matmul(U, V)
return V
def __init__(self, shape, n_components, a=0.1, b=1):
tl.set_backend('numpy')
self.shape = shape
self.n_components = n_components
self.n_modes = len(shape)
self.Z = [] # factors
self.X = None # tensor
# Gamma shape(A) and mean(B) objects. Gamma scale is B/A
self.A = self._create_gamma_prior(a)
self.B = self._create_gamma_prior(b)
# Inference variables:
self.C = [] # variational shape parameters for factors
self.D = [] # variational scale parameters for factors
self.E = [] # arithmetic expectations of factors
self.L = [] # geometric expectations of factors
# Randomly initialize itself
self.rand_init()
def create_cp_sparse_gen(dims, rank, n_el, method='rand', return_sparse=False):
tl.set_backend('pytorch')
if method == 'rand':
randfunc = torch.rand
elif method == 'randn':
randfunc = torch.randn
else:
raise NotImplementedError(f'Unknown random method: {method}')
n_dims = len(dims)
factors = [randfunc((dim, rank)) for dim in dims]
lambdas = torch.ones(rank)
# Create probability tensor
P = normalize_cp_ten((lambdas, factors))
lambdas /= torch.sum(lambdas)
# Count samples per component
n_edges = n_el
if n_edges < 1:
n_edges = round(n_edges * torch.prod(dims))
csums = probsample(n_edges, lambdas)
subs = []
# Calculate subscripts
for c in range(rank):
n_sample = int(csums[c])
if n_sample == 0:
continue
sub_idxs = torch.zeros((int(n_sample), n_dims), dtype=torch.int64)
for d in range(n_dims):
sub_idxs[:, d] = probsample(n_sample,
factors[d][:, c],
return_counts=False)
subs.append(sub_idxs)
all_subs = torch.vstack(subs).T
sp_ten = torch.sparse_coo_tensor(all_subs, torch.ones(n_edges))
if return_sparse:
return sp_ten
return sp_ten.to_dense()
def main():
tl.set_backend("pytorch")
args = parser.parse_args()
if args.rank is not None and args.threshold is not None:
raise Exception(
"Conflicting arguments passed: args.rank and args.threshold.\n\tYou can only set either rank argument or threshold argument. You can't set both."
)
if args.resume is not None and args.reset_weights:
raise Exception(
"Conflicting arguments passed: args.resume and args.reset_weights.\n\tYou can't resume training from a certain checkpoint and reset weights as well."
)
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
if args.gpu is not None:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
if args.dist_url == "env://" and args.world_size == -1:
args.world_size = int(os.environ["WORLD_SIZE"])
args.distributed = args.world_size > 1 or args.multiprocessing_distributed
ngpus_per_node = torch.cuda.device_count()
if args.multiprocessing_distributed:
# Since we have ngpus_per_node processes per node, the total world_size
# needs to be adjusted accordingly
args.world_size = ngpus_per_node * args.world_size
# Use torch.multiprocessing.spawn to launch distributed processes: the
# main_worker process function
mp.spawn(main_worker,
nprocs=ngpus_per_node,
args=(ngpus_per_node, args))
else:
# Simply call main_worker function
main_worker(args.gpu, ngpus_per_node, args)
def cp_decomposition_net(net):
tl.set_backend('pytorch')
N = len(net.features._modules.keys())
for i, key in enumerate(net.features._modules.keys()):
if i >= N - 2:
break
if isinstance(net.features._modules[key], torch.nn.modules.conv.Conv2d):
conv_layer = net.features._modules[key]
rank = max(conv_layer.weight.data.cpu().numpy().shape) // 3
decomposed = cp_decomposition_conv_layer(conv_layer, rank) # CP分解
net.features._modules[key] = decomposed
for param in net.parameters():
param.requires_grad = True
if torch.cuda.is_available():
net.cuda()
return net
def __init__(self, shape, n_components, alpha=0.1, beta=1):
tl.set_backend('numpy')
self.shape = shape
self.n_components = n_components
self.n_modes = len(shape)
self.T = None # factor 1
self.V = None # factor 2
self.X = None # tensor
# Gamma shape(A) and mean(B) objects. Gamma scale is B/A
self.At = self.Av = None
self.Bt = self.Bv = None
# Inference variables:
self.Ct = self.Cv = None
self.Dt = self.Dv = None
self.Et = self.Ev = None
self.Lt = self.Lv = None
# Randomly initialize itself
self.rand_init(alpha, beta)
self.rand_init(alpha, beta)
def run_parafac(path,
dimensions=20,
nonnegative=False,
is_matrix=True,
name='current',
cuda=False,
noise='orig',
iterations=1,
fixed='False'):
# setting up the backend
if cuda:
tl.set_backend('pytorch')
# checking if the arguments are valid
assert is_ndarray_folder(path), "Invalid path."
assert '_' not in name, "Name cannot contain the '_' symbol."
# creating some useful paths to store factorization results
mat_path = join(
path, 'mat_' + str(dimensions) + '_' + str(iterations) + '_' + name)
ten_path = join(
path, 'ten_' + str(dimensions) + '_' + str(iterations) + '_' + name)
# loading the meta data
with open(join(path, 'meta_data.json'), 'r') as json_file:
meta_data = json.load(json_file)
# removing old factorization with same name (if exists)
delete_factorization_by_name(name, path)
# factorizing the data
start = time.time()
if is_matrix:
matrix = sparse.load_npz(join(path, 'matrix.npz')).todense()
matrices = prepare_ndarrays(matrix, iterations, fixed, noise)
create_folder_if_absent(mat_path)
factorize_matrices(matrices, iterations, dimensions, nonnegative,
mat_path, cuda)
else:
tensor = sparse.load_npz(join(path, 'tensor.npz')).todense()
tensors = prepare_ndarrays(tensor, iterations, fixed, noise)
create_folder_if_absent(ten_path)
factorize_tensors(tensors, iterations, dimensions, nonnegative,
ten_path, meta_data, path, cuda)
end = time.time()
print('Factorization completed in %d seconds' % (end - start))
def build(model, decomp='cp'):
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
print('==> Building model..')
tl.set_backend('pytorch')
full_net = model
full_net = full_net.to(device)
path = 'models/'
if not os.path.exists(path):
os.mkdir(path)
torch.save(full_net, path + 'model')
if decomp:
decompose_conv(decomp)
decompose_fc(decomp)
if device == 'cuda:0':
net = torch.load(path + "model").cuda()
else:
net = torch.load(path + "model")
print(net)
print('==> Done')
return net
def get_dataset(kde, **option):
option_default = {
'env': gym.make('Pendulum-v0'),
'num_trajs': 1000,
'max_episode_length': 10,
'window_size': 5
}
option = {**option_default, **option}
action_all_pr, observation_all_pr, x_pr, new_rewards_pr = generate_data(
**option)
pr = construct_PR_target(kde, x_pr, new_rewards_pr)
tl.set_backend('pytorch')
pr = normalize(pr)
new_data = Dataset(data=[
tl.tensor(action_all_pr).float(),
tl.tensor(observation_all_pr).float(),
tl.tensor(pr).float()
])
return new_data
def __init__(self,
mask=None,
tol=10e-7,
reg_E=1,
reg_J=1,
mu_init=10e-5,
mu_max=10e9,
learning_rate=1.1,
n_iter_max=100,
verbose=1,
backend='numpy'):
self.mask = mask
self.tol = tol
self.reg_E = reg_E
self.reg_J = reg_J
self.mu_init = mu_init
self.mu_max = mu_max
self.lr = learning_rate
self.n_iter_max = n_iter_max
self.verbose = verbose
self.backend = backend
tl.set_backend(backend)
def get_cp_factors(self):
tl.set_backend('pytorch')
bias = self.bias
if isinstance(self.layer, nn.Sequential):
# Tensorly case
_, (f_cout, f_cin, f_h, f_w) = parafac(kruskal_to_tensor(
(None, self.weight)),
self.rank,
n_iter_max=5000,
init='random',
tol=1e-8,
svd=None,
cvg_criterion='rec_error')
else:
# Tensorly case
_, (f_cout, f_cin, f_h, f_w) = parafac(self.weight,
self.rank,
n_iter_max=5000,
init='random',
tol=1e-8,
svd=None,
cvg_criterion='rec_error')
# # Reshape factor matrices to 4D weight tensors
# f_cin: (cin, rank) -> (rank, cin, 1, 1)
# f_h: (h, rank) -> (rank, 1, h, 1)
# f_w: (w, rank) -> (rank, 1, 1, w)
# f_cout: (count, rank) -> (count, rank, 1, 1)
# Pytorh case
f_cin = f_cin.t().unsqueeze_(2).unsqueeze_(3).contiguous()
f_h = f_h.t().unsqueeze_(1).unsqueeze_(3).contiguous()
f_w = f_w.t().unsqueeze_(1).unsqueeze_(2).contiguous()
f_cout = f_cout.unsqueeze_(2).unsqueeze_(3).contiguous()
return [f_cin, f_h, f_w, f_cout], [None, None, None, bias]
parser.add_argument("--decompose", dest="decompose", action="store_true")
parser.add_argument("--fine_tune", dest="fine_tune", action="store_true")
parser.add_argument("--train_path", type = str, default = "train")
parser.add_argument("--test_path", type = str, default = "test")
parser.add_argument("--cp", dest="cp", action="store_true", \
help="Use cp decomposition. uses tucker by default")
parser.set_defaults(train=False)
parser.set_defaults(decompose=False)
parser.set_defaults(fine_tune=False)
parser.set_defaults(cp=False)
args = parser.parse_args()
return args
if __name__ == '__main__':
args = get_args()
tl.set_backend('numpy')
if args.train:
model = ModifiedVGG16Model().cuda()
optimizer = optim.SGD(model.classifier.parameters(), lr=0.0001, momentum=0.99)
trainer = Trainer(args.train_path, args.test_path, model, optimizer)
trainer.train(epoches = 10)
torch.save(model, "model")
elif args.decompose:
model = torch.load("model").cuda()
model.eval()
model.cpu()
N = len(model.features._modules.keys())
for i, key in enumerate(model.features._modules.keys()):
