from vit_jax import checkpoint from vit_jax import input_pipeline from vit_jax import utils from vit_jax import models from vit_jax import train from vit_jax.configs import common as common_config from vit_jax.configs import models as models_config from absl import logging import flax import jax from matplotlib import pyplot as plt import numpy as np import optax import tqdm import os logging.set_verbosity(logging.INFO) import PIL import tensorflow_datasets as tfds import time # import tensorflow as tf '''测试dcu/gpu''' from jax.lib import xla_bridge jax_test=xla_bridge.get_backend().platform if not (jax_test=='gpu'): exit() '''指定模型''' model_name = 'ViT-B_16' #@param ["ViT-B_32", "Mixer-B_16"] pretrained_path=f'./test_result/{model_name}.npz' model_path=f'./test_result/{model_name}_imagenet2012.npz' '''加载数据集--微调用''' # dataset = 'cifar100' # imagenet2012 cifar10 cifar100 # batch_size = 512 # config = common_config.with_dataset(common_config.get_config(), dataset) # # config.shuffle_buffer=1000 # # config.accum_steps=64 # config.batch = batch_size # config.pp.crop = 384 # # 建立数据集 # ds_train = input_pipeline.get_data_from_tfds(config=config, mode='train') # ds_test = input_pipeline.get_data_from_tfds(config=config, mode='test') # num_classes = input_pipeline.get_dataset_info(dataset, 'train')['num_classes'] # del config # Only needed to instantiate datasets. # # Fetch a batch of test images for illustration purposes. # batch = next(iter(ds_test.as_numpy_iterator())) # # Note the shape : [num_local_devices, local_batch_size, h, w, c] # print("数据集shape：",batch['image'].shape) '''加载预训练模型--微调用''' # model_config = models_config.MODEL_CONFIGS[model_name] # print("模型config：",model_config) # # 加载模型定义并初始化随机参数。 # # 这也将模型编译为XLA(第一次需要几分钟)。 # if model_name.startswith('Mixer'): # model = models.MlpMixer(num_classes=num_classes, **model_config) # else: # model = models.VisionTransformer(num_classes=num_classes, **model_config) # variables = jax.jit(lambda: model.init( # jax.random.PRNGKey(0), # # 丢弃用于初始化的批处理的“num_local_devices”维度。 # batch['image'][0, :1], # train=False, # ), backend='cpu')() # #加载和转换预训练检查点。 # # 这涉及到加载实际的预训练模型结果，但也要修改一点参数，例如改变最终层，并调整位置嵌入的大小。有关详细信息，请参阅代码和本文的方法。 # params = checkpoint.load_pretrained( # pretrained_path=pretrained_path, # init_params=variables['params'], # model_config=model_config # ) '''评估''' # params_repl = flax.jax_utils.replicate(params) # print('params.cls:', type(params['head']['bias']).__name__, # params['head']['bias'].shape) # print('params_repl.cls:', type(params_repl['head']['bias']).__name__, # params_repl['head']['bias'].shape) # # 然后将调用映射到我们模型的forward pass到所有可用的设备。 # vit_apply_repl = jax.pmap(lambda params, inputs: model.apply( # dict(params=params), inputs, train=False)) # def get_accuracy(params_repl): # """返回对测试集求值的精度""" # good = total = 0 # steps = input_pipeline.get_dataset_info(dataset, 'test')['num_examples'] // batch_size # for _, batch in zip(tqdm.trange(steps), ds_test.as_numpy_iterator()): # predicted = vit_apply_repl(params_repl, batch['image']) # is_same = predicted.argmax(axis=-1) == batch['label'].argmax(axis=-1) # good += is_same.sum() # total += len(is_same.flatten()) # return good / total # # 模型的随机性能 # print(get_accuracy(params_repl)) '''微调''' # # 100 Steps take approximately 15 minutes in the TPU runtime. # total_steps = 50 # warmup_steps = 5 # decay_type = 'cosine' # grad_norm_clip = 1 # # 这控制了批处理被分割的转发次数。8适用于具有8个设备的TPU运行时。64应该可以在GPU上工作。当然，您也可以调整上面的batch_size，但这需要相应地调整学习率。 # accum_steps = 64 # TODO:可能要改 # base_lr = 0.03 # # 检查 train.make_update_fn # lr_fn = utils.create_learning_rate_schedule(total_steps, base_lr, decay_type, warmup_steps) # # 我们使用一个动量优化器，使用一半精度的状态来节省内存。它还实现了梯度裁剪 # tx = optax.chain( # optax.clip_by_global_norm(grad_norm_clip), # optax.sgd( # learning_rate=lr_fn, # momentum=0.9, # accumulator_dtype='bfloat16', # ), # ) # update_fn_repl = train.make_update_fn( # apply_fn=model.apply, accum_steps=accum_steps, tx=tx) # opt_state = tx.init(params) # opt_state_repl = flax.jax_utils.replicate(opt_state) # # Initialize PRNGs for dropout. # update_rng_repl = flax.jax_utils.replicate(jax.random.PRNGKey(0)) # # 训练更新 # losses = [] # lrs = [] # # Completes in ~20 min on the TPU runtime. # start = time.time() # for step, batch in zip( # tqdm.trange(1, total_steps + 1), # ds_train.as_numpy_iterator(), # ): # params_repl, opt_state_repl, loss_repl, update_rng_repl = update_fn_repl( # params_repl, opt_state_repl, batch, update_rng_repl) # losses.append(loss_repl[0]) # lrs.append(lr_fn(step)) # end = time.time() # print(f"{model_name}_{dataset}_{total_steps}_{warmup_steps}微调时间为：",end-start) # print(get_accuracy(params_repl)) # 绘制学习率变化曲线并保存 # plt.plot(losses) # plt.savefig(f'./test_result/{model_name}_{dataset}/losses_plot.png') # plt.close() # plt.plot(lrs) # plt.savefig(f'./test_result/{model_name}_{dataset}/lrs_plot.png') # plt.close() # 在CIFAR10上，Mixer-B/16应该是~96.7%，vitb /32应该是97.7%(都是@224) '''推理''' model_config = models_config.MODEL_CONFIGS[model_name] print("模型config：",model_config) model = models.VisionTransformer(num_classes=1000, **model_config) assert os.path.exists(model_path) # 加载和转换预训练的检查点 params = checkpoint.load(model_path) params['pre_logits'] = {} # Need to restore empty leaf for Flax. # 获取图像标签. # get_ipython().system('wget https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt') imagenet_labels = dict(enumerate(open('./dataset/ilsvrc2012_wordnet_lemmas.txt'))) # 得到一张具有正确尺寸的随机图片 # resolution = 224 if model_name.startswith('Mixer') else 384 # get_ipython().system('wget https://picsum.photos/$resolution -O picsum.jpg') img = PIL.Image.open('./dataset/picsum.jpg') # 预测 start_time=time.time() logits, = model.apply(dict(params=params), (np.array(img) / 128 - 1)[None, ...], train=False) end_time=time.time() preds = np.array(jax.nn.softmax(logits)) print("推理结果：time=",end_time-start_time) for idx in preds.argsort()[:-11:-1]: print(f'{preds[idx]:.5f} : {imagenet_labels[idx]}', end='')