Bò sát: Một thuật toán siêu học có thể mở rộng

importnumpyasnp

importtorch

, fromtorchimportnnautogradasag

. importmatplotlibpyplotasplt

fromcopyimportdeepcopy



seed=0

plot=True

innerstepsize=0.02# stepsize in inner SGD

innerepochs=1# number of epochs of each inner SGD

outerstepsize0=0.1# stepsize of outer optimization, i.e., meta-optimization

niterations=30000# number of outer updates; each iteration we sample one task and update on it



..()rng=nprandomRandomStateseed

.()torchmanual_seedseed



# Define task distribution

.(, , )[:,] x_all=nplinspace-5550None# All of the x points

ntrain=10# Size of training minibatches

():defgen_task

"Generate classification problem"

.(, .)phase=rnguniformlow=0high=2*nppi

.(, )ampl=rnguniform0.15

: .( ) f_randomsine=lambdaxnpsinx+phase*ampl

returnf_randomsine



# Define model. Reptile paper uses ReLU, but Tanh gives slightly better results

.(model=nnSequential

.(, ),nnLinear164

.(),nnTanh

.(, ),nnLinear6464

.(),nnTanh

.(, ),nnLinear641

)



():deftotorchx

.(.())returnagVariabletorchTensorx



(, ):deftrain_on_batchxy

()x=totorchx

()y=totorchy

.()modelzero_grad

()ypred=modelx

( ).().()loss=ypred-ypow2mean

.()lossbackward

.():forparaminmodelparameters

. ..paramdata-=innerstepsize*paramgraddata



():defpredictx

()x=totorchx

()..()returnmodelxdatanumpy



# Choose a fixed task and minibatch for visualization

()f_plot=gen_task

[.((), )]xtrain_plot=x_allrngchoicelenx_allsize=ntrain



# Reptile training loop

():foriterationinrangeniterations

(.())weights_before=deepcopymodelstate_dict

# Generate task

()f=gen_task

()y_all=fx_all

# Do SGD on this task

.(())inds=rngpermutationlenx_all

():for_inrangeinnerepochs

(, (), ):forstartinrange0lenx_allntrain

[:]mbinds=indsstartstart+ntrain

([], [])train_on_batchx_allmbindsy_allmbinds

# Interpolate between current weights and trained weights from this task

# I.e. (weights_before - weights_after) is the meta-gradient

.()weights_after=modelstate_dict

( ) outerstepsize=outerstepsize0*1-iteration/niterations# linear schedule

.({ :modelload_state_dictname

[] ([] []) weights_beforename+weights_aftername-weights_beforename*outerstepsize

})fornameinweights_before



# Periodically plot the results on a particular task and minibatch

() :ifplotanditeration==0oriteration+1%1000==0

.()pltcla

f=f_plot

(.()) weights_before=deepcopymodelstate_dict# save snapshot before evaluation

.(, (), , (,,))pltplotx_allpredictx_alllabel="pred after 0"color=001

():forinneriterinrange32

(, ())train_on_batchxtrain_plotfxtrain_plot

() :ifinneriter+1%8==0

() frac=inneriter+1/32

.(, (), (), (, , ))pltplotx_allpredictx_alllabel="pred after %i"%inneriter+1color=frac01-frac

.(, (), , (,,))pltplotx_allfx_alllabel="true"color=010

.(() ()).()lossval=npsquarepredictx_all-fx_allmean

.(, (), , , )pltplotxtrain_plotfxtrain_plot"x"label="train"color="k"

.(,)pltylim-44

.()pltlegendloc="lower right"

.()pltpause0.01

.() modelload_state_dictweights_before# restore from snapshot

()printf"-----------------------------"

()printf"iteration "{iteration+1}

() printf"loss on plotted curve ":.3f{lossval}# would be better to average loss over a set of examples, but this is optimized for brevity