搜索

Tensorflow 使用TPU训练


发布时间: 2022-11-24 17:57:03    浏览次数:8 次

  要用TPU训练tensorflow模型,只能使用静态图。也就是要先通过keras的sequential或者函数式定义模型,而不能直接使用重写的Model类。例子如下,其中包含层的自定义,以及子像素卷积。需要注意的是,tensorflow的子pixel_shuffle通道顺序与pytorch不同,具体怎么不同不记录了,可以直接实验一下。

from tensorflow import keras
from tensorflow.keras import losses,layers,optimizers,Model
import tensorflow as tf
import numpy as np 

tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.experimental.TPUStrategy(tpu) 

def pixel_unshuffle(x, scale):
  x = tf.nn.space_to_depth(x, scale)
  return x

class MyDense(layers.Layer): 
  def __init__(self): 
    super().__init__()   
    self.layer = layers.Conv2D(3, 3, 1, 'same')
  def call(self, inp): 
    x = self.layer(inp)
    x = pixel_unshuffle(x, 2)
    x = tf.maximum(x, 50)
    return x

with strategy.scope():
  inputs = keras.Input(shape=[48,48,3]) 
  x = MyDense()(inputs)
  model = Model(inputs, x)
  model.compile(optimizers.SGD(), losses.MSE)

x = np.zeros([4096*10,48,48,3]).astype(np.float32)
y = np.zeros([4096*10,24,24,12]).astype(np.float32)
model.fit(x,y,epochs=50,batch_size=4096)

 

免责声明 Tensorflow 使用TPU训练,资源类别:文本, 浏览次数:8 次, 文件大小:-- , 由本站蜘蛛搜索收录2022-11-24 05:57:03。此页面由程序自动采集,只作交流和学习使用,本站不储存任何资源文件,如有侵权内容请联系我们举报删除, 感谢您对本站的支持。 原文链接:https://www.cnblogs.com/qizhou/p/14873096.html