我正在Pytorch中训练模型,我想使用输入的截断SVD分解。为了计算SVD,我将输入的是Pytorch Cuda张量传输到CPU,然后使用TruncatedSVD
中的scikit-learn
执行截断,然后将结果传输回GPU。以下是我的模型的代码:
class ImgEmb(nn.Module):
def __init__(self, input_size, hidden_size):
super(ImgEmb, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.drop = nn.Dropout(0.2)
self.mlp = nn.Linear(input_size/2, hidden_size)
self.relu = nn.Tanh()
self.svd = TruncatedSVD(n_components=input_size/2)
def forward(self, input):
svd=self.svd.fit_transform(input.cpu())
svd_tensor=torch.from_numpy(svd)
svd_tensor=svd_tensor.cuda()
mlp=self.mlp(svd_tensor)
res = self.relu(mlp)
return res
我想知道是否可以在不往返于GPU的情况下实现截断的SVD吗? (因为这很耗时并且根本没有效率)
[您可以直接使用PyTorch的SVD并手动截断它,也可以将TensorLy中的截短的SVD与PyTorch后端一起使用: