问题描述
我正在尝试使用Keras在google Colab上实现自我注意GAN。当我测试注意层时，出现OOM错误。那么，我在矩阵乘法上做错什么了吗？或者对于高分辨率（> 64 x 64）的colab GPU来说，这仅仅是一个昂贵的操作？
def hw_flatten(x):
# Input shape x: [BATCH,HEIGHT,WIDTH,CHANNELS]
# flat the feature volume across the width and height dimensions
x = Reshape((x.shape[1]*x.shape[2],x.shape[3]))(x) #in the Reshape layer batch is implicit
return x # return [BATCH,W*H,CHANNELS]
def matmul(couple_t):
tensor_1 = couple_t[0]
tensor_2 = couple_t[1]
transponse = couple_t[2] #boolean
return tf.matmul(tensor_1,tensor_2,transpose_b=transponse)
class SelfAttention(Layer):
def _init_(self,ch,**kwargs):
super(SelfAttention,self).__init__(**kwargs)
self.ch = ch
def attentionMap(self,feature_map):
f = Conv2D(filters=feature_map.shape[3]/8,kernel_size=(1,1),strides=1,padding='same')(feature_map) # [bs,h,w,c']
g = Conv2D(filters=feature_map.shape[3]/8,c']
h = Conv2D(filters=feature_map.shape[3],padding='same')(feature_map) # [bs,c']
s = Lambda(matmul)([hw_flatten(g),hw_flatten(f),True]) # [bs,N,N]
beta = Activation("softmax")(s)
o = Lambda(matmul)([beta,hw_flatten(h),False]) # [bs,C]
gamma = self.add_weight(name='gamma',shape=[1],initializer='zeros',trainable=True)
o = Reshape((feature_map.shape[1:]))(o) # [bs,C]
x = gamma * o + feature_map
print(x.shape)
return x
这是测试：
tensor = np.random.normal(0,1,size=(32,64,512)).astype('float64')
attention_o = SelfAttention(64)
a = attention_o.attentionMap(tensor)
这是错误：
OOM when allocating tensor with shape[32,4096,4096] and type double
非常感谢您的关注：D
解决方法
暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！
如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。
小编邮箱:dio#foxmail.com (将#修改为@）