Google Colabで TPUを使いたい

前提・実現したいこと

お世話になります
機械学習初心者ですよろしくお願い致します

Google ColaboratoryのTPUが遅い を参考にさせていただき
こちらを (pix2pix のプログラムです) TPUを使用して実行させたいと思っております

発生している問題・エラーメッセージ

python
1Traceback (most recent call last):
2  File "pix2pix-TPU.py", line 146, in <module>
3    training_updates = tf.train.AdamOptimizer(learning_rate=lrD, beta1=0.5).get_updates(netD.trainable_weights,[],loss_D)
4AttributeError: 'AdamOptimizer' object has no attribute 'get_updates'

と表示されます

機械学習を　始めて間もないため
tensorflowが私には難しく kerasからの学習をしております
kerasでの 'get_updates'をtensorflowで
どう記述してよいのかが皆目解りません

どうぞよろしくお願い致します

試したこと

TPUの恩恵を受けられるは別として
146、151行をコメントアウトして 145,150行をアンコメントすると

python
1Traceback (most recent call last):
2  File "pix2pix-TPU.py", line 145, in <module>
3    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(netD.trainable_weights,[],loss_D)
4TypeError: get_updates() takes 3 positional arguments but 4 were given

となります

補足情報（FW/ツールのバージョンなど）

改変コードが字数制限で掲載できません

python
1import os
2
3import tensorflow as tf
4import tensorflow.keras.backend as K 
5
6from tensorflow.contrib.tpu.python.tpu import keras_support
7K.set_image_data_format('channels_last')
8channel_axis=-1
9channel_first = False
10
11
12from tensorflow.keras.models import Sequential, Model
13from tensorflow.keras.layers import Conv2D, ZeroPadding2D, BatchNormalization, Input, Dropout
14from tensorflow.keras.layers import Conv2DTranspose, Reshape, Activation, Cropping2D, Flatten
15from tensorflow.keras.layers import Concatenate
16from tensorflow.keras.layers import LeakyReLU
17from tensorflow.keras.activations import relu
18from tensorflow.keras.initializers import RandomNormal
19
20
21def __conv_init(a):
22    print("conv_init", a)
23    k = RandomNormal(0, 0.02)(a) 
24    k.conv_weight = True    
25    return k
26conv_init = RandomNormal(0, 0.02)
27gamma_init = RandomNormal(1., 0.02) 
28
29
30def conv2d(f, *a, **k):
31    return Conv2D(f, kernel_initializer = conv_init, *a, **k)
32def batchnorm():
33    return BatchNormalization(momentum=0.9, axis=channel_axis, epsilon=1.01e-5,
34                                   gamma_initializer = gamma_init)
35def BASIC_D(nc_in, nc_out, ndf, max_layers=3):
36   
37    if channel_first:
38        input_a, input_b =  Input(shape=(nc_in, None, None)), Input(shape=(nc_out, None, None))
39    else:
40        input_a, input_b = Input(shape=(None, None, nc_in)), Input(shape=(None, None, nc_out))
41    _ = Concatenate(axis=channel_axis)([input_a, input_b])
42    _ = conv2d(ndf, kernel_size=4, strides=2, padding="same", name = 'First') (_)
43    _ = LeakyReLU(alpha=0.2)(_)
44    
45    for layer in range(1, max_layers):        
46        out_feat = ndf * min(2**layer, 8)
47        _ = conv2d(out_feat, kernel_size=4, strides=2, padding="same", 
48                   use_bias=False, name = 'pyramid.{0}'.format(layer)             
49                        ) (_)
50        _ = batchnorm()(_, training=1)        
51        _ = LeakyReLU(alpha=0.2)(_)
52    
53    out_feat = ndf*min(2**max_layers, 8)
54    _ = ZeroPadding2D(1)(_)
55    _ = conv2d(out_feat, kernel_size=4,  use_bias=False, name = 'pyramid_last') (_)
56    _ = batchnorm()(_, training=1)
57    _ = LeakyReLU(alpha=0.2)(_)
58    
59    # final layer
60    _ = ZeroPadding2D(1)(_)
61    _ = conv2d(1, kernel_size=4, name = 'final'.format(out_feat, 1), 
62               activation = "sigmoid") (_)    
63    return Model(inputs=[input_a, input_b], outputs=_)
64
65def UNET_G(isize, nc_in=3, nc_out=3, ngf=64, fixed_input_size=True):    
66    max_nf = 8*ngf    
67    def block(x, s, nf_in, use_batchnorm=True, nf_out=None, nf_next=None):
68       
69        assert s>=2 and s%2==0
70        if nf_next is None:
71            nf_next = min(nf_in*2, max_nf)
72        if nf_out is None:
73            nf_out = nf_in
74        x = conv2d(nf_next, kernel_size=4, strides=2, use_bias=(not (use_batchnorm and s>2)),
75                   padding="same", name = 'conv_{0}'.format(s)) (x)
76        if s>2:
77            if use_batchnorm:
78                x = batchnorm()(x, training=1)
79            x2 = LeakyReLU(alpha=0.2)(x)
80            x2 = block(x2, s//2, nf_next)
81            x = Concatenate(axis=channel_axis)([x, x2])            
82        x = Activation("relu")(x)
83        x = Conv2DTranspose(nf_out, kernel_size=4, strides=2, use_bias=not use_batchnorm,
84                            kernel_initializer = conv_init,          
85                            name = 'convt.{0}'.format(s))(x)        
86        x = Cropping2D(1)(x)
87        if use_batchnorm:
88            x = batchnorm()(x, training=1)
89        if s <=8:
90            x = Dropout(0.5)(x, training=1)
91        return x
92    
93    s = isize if fixed_input_size else None
94    if channel_first:
95        _ = inputs = Input(shape=(nc_in, s, s))
96    else:
97        _ = inputs = Input(shape=(s, s, nc_in))        
98    _ = block(_, isize, nc_in, False, nf_out=nc_out, nf_next=ngf)
99    _ = Activation('tanh')(_)
100    return Model(inputs=inputs, outputs=[_])
101
102nc_in = 3
103nc_out = 3
104ngf = 64
105ndf = 64
106λ = 10
107
108loadSize = 256
109imageSize = 256
110batchSize = 1
111lrD = 2e-4
112lrG = 2e-4
113
114
115netD = BASIC_D(nc_in, nc_out, ndf)
116netD.summary()
117
118from IPython.display import SVG
119
120netG = UNET_G(imageSize, nc_in, nc_out, ngf)
121netG.summary()
122
123
124
125
126from tensorflow.keras.optimizers import RMSprop, SGD, Adam
127
128real_A = netG.input
129fake_B = netG.output
130netG_generate = K.function([real_A], [fake_B])
131real_B = netD.inputs[1]
132output_D_real = netD([real_A, real_B])
133output_D_fake = netD([real_A, fake_B])
134
135loss_fn = lambda output, target : -K.mean(K.log(output+1e-12)*target+K.log(1-output+1e-12)*(1-target))
136
137loss_D_real = loss_fn(output_D_real, K.ones_like(output_D_real))
138loss_D_fake = loss_fn(output_D_fake, K.zeros_like(output_D_fake))
139loss_G_fake = loss_fn(output_D_fake, K.ones_like(output_D_fake))
140
141
142loss_L1 = K.mean(K.abs(fake_B-real_B))
143
144loss_D = loss_D_real +loss_D_fake
145training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(netD.trainable_weights,[],loss_D)
146#training_updates = tf.train.AdamOptimizer(learning_rate=lrD, beta1=0.5).get_updates(netD.trainable_weights,[],loss_D)
147netD_train = K.function([real_A, real_B],[loss_D/2], training_updates)
148
149loss_G = loss_G_fake   + 100 * loss_L1
150training_updates = Adam(lr=lrG, beta_1=0.5).get_updates(netG.trainable_weights,[], loss_G)
151#training_updates = tf.train.AdamOptimizer(learning_rate=lrG, beta1=0.5).get_updates(netG.trainable_weights,[], loss_G)
152netG_train = K.function([real_A, real_B], [loss_G_fake, loss_L1], training_updates)
153
154try:
155    netG.load_weights('netG.hdf5')
156    netD.load_weights('netD.hdf5')
157    print('重みデータの読み込み完了')
158except:
159    print('重みデータを読み込めませんでした')
160    pass
161
162tpu_grpc_url = "grpc://"+os.environ["COLAB_TPU_ADDR"]
163tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
164strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
165netD = tf.contrib.tpu.keras_to_tpu_model(netD, strategy=strategy)
166netG = tf.contrib.tpu.keras_to_tpu_model(netG, strategy=strategy)