|
- # Copyright 2022 Huawei Technologies Co., Ltd
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- # ============================================================================
-
- """export file of MINDIR format"""
-
- import argparse
- import numpy as np
- import math
- import time, random
- import mindspore as ms
- from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
-
- from model import load_decoder_arch, load_encoder_arch
-
- def get_args():
- parser = argparse.ArgumentParser(description='CFLOW-AD')
- parser.add_argument('--dataset', default='mvtec', type=str, metavar='D',
- help='dataset name: mvtec/stc (default: mvtec)')
- parser.add_argument('--device_target', default='Ascend', help='enables npu')
- parser.add_argument('--device_id', type=int, default=0)
- parser.add_argument('--checkpoint_dir', default='', type=str, metavar='D',
- help='file with saved checkpoint')
- parser.add_argument('--save_checkpoint_dir', default='./weights', type=str, metavar='D',
- help='file with save checkpoint')
- parser.add_argument('--encoder_ckpt_pre_dir', default='epoch_22_mvtec_wide_resnet50_2_freia-cflow_pl3_cb8_inp256_run0_cable_2022-10-06-21_02_19_encoder', type=str, metavar='D',
- help='pre file name with saved checkpoint')
- parser.add_argument('--decoder_ckpt_pre_dir', default='epoch_22_mvtec_wide_resnet50_2_freia-cflow_pl3_cb8_inp256_run0_cable_2022-10-06-21_02_19_decoder', type=str, metavar='D',
- help='pre file name with saved checkpoint')
- parser.add_argument('-cl', '--class-name', default='none', type=str, metavar='C',
- help='class name for MVTec/STC (default: none)')
- parser.add_argument('-enc', '--enc-arch', default='wide_resnet50_2', type=str, metavar='A',
- help='feature extractor: wide_resnet50_2/resnet18/mobilenet_v3_large (default: wide_resnet50_2)')
- parser.add_argument('-dec', '--dec-arch', default='freia-cflow', type=str, metavar='A',
- help='normalizing flow model (default: freia-cflow)')
- parser.add_argument('-pl', '--pool-layers', default=3, type=int, metavar='L',
- help='number of layers used in NF model (default: 3)')
- parser.add_argument('-cb', '--coupling-blocks', default=8, type=int, metavar='L',
- help='number of layers used in NF model (default: 8)')
- parser.add_argument('-run', '--run-name', default=0, type=int, metavar='C',
- help='name of the run (default: 0)')
- parser.add_argument('-inp', '--input-size', default=256, type=int, metavar='C',
- help='image resize dimensions (default: 256)')
- parser.add_argument("--action-type", default='norm-train', type=str, metavar='T',
- help='norm-train/norm-test (default: norm-train)')
- parser.add_argument('-bs', '--batch-size', default=32, type=int, metavar='B',
- help='train batch size (default: 32)')
- parser.add_argument('--lr', type=float, default=1e-3, metavar='LR',
- help='learning rate (default: 2e-4)')
- parser.add_argument('--meta-epochs', type=int, default=25, metavar='N',
- help='number of meta epochs to train (default: 25)')
- parser.add_argument('--sub-epochs', type=int, default=8, metavar='N',
- help='number of sub epochs to train (default: 8)')
- parser.add_argument('--pro', action='store_true', default=False,
- help='enables estimation of AUPRO metric')
- parser.add_argument('--viz', action='store_true', default=False,
- help='saves test data visualizations')
- parser.add_argument('--workers', default=4, type=int, metavar='G',
- help='number of data loading workers (default: 4)')
- parser.add_argument("--gpu", default='0', type=str, metavar='G',
- help='GPU device number')
- parser.add_argument('--no-cuda', action='store_true', default=False,
- help='disables CUDA training')
- parser.add_argument('--video-path', default='.', type=str, metavar='D',
- help='video file path')
- parser.add_argument("--L", type=int, default=3, help="While exporting encoder, L means the number of layers.")
- parser.add_argument("--exported_ckpt_path", type=str, default="./ckpt/wide_resnet50_2_", help="Prefix of checkpoint file path.")
- parser.add_argument("--exported_file_name", type=str, default="CFLOW-AD", help="output file name.")
- parser.add_argument("--export_mode", type=str, default="decoder", help="Choose export 'encoder' or 'decoder'.")
- parser.add_argument("--file_format", type=str, default="MINDIR", help="output file format")
- parser.add_argument("--data_url", type=str, default="", help="dataset path")
-
- args = parser.parse_args()
- return args
-
- def init_seeds(seed=0):
- random.seed(seed)
- np.random.seed(seed)
- ms.set_seed(seed)
-
-
- if __name__ == '__main__':
- #environment
- c = get_args()
- c.condition_vec = 128
- c.img_size = (c.input_size, c.input_size)
- c.img_dims = [3] + list(c.img_size)
- # network hyperparameters
- c.clamp_alpha = 1.9 # see paper equation 2 for explanation
- c.dropout = 0.0 # dropout in s-t-networks
- c.bs=c.batch_size
- # dataloader parameters
- if c.dataset == 'mvtec':
- c.data_path = c.data_url
- c.verbose = True
- c.hide_tqdm_bar = True
- c.save_results = True
- # unsup-train
- c.print_freq = 2
- c.temp = 0.5
- c.lr_decay_epochs = [i*c.meta_epochs//100 for i in [50,75,90]]
- c.lr_decay_rate = 0.1
- c.lr_warm_epochs = 2
- c.lr_warm = True
- c.lr_cosine = True
- if c.lr_warm:
- c.lr_warmup_from = c.lr/10.0
- if c.lr_cosine:
- eta_min = c.lr * (c.lr_decay_rate ** 3)
- c.lr_warmup_to = eta_min + (c.lr - eta_min) * (
- 1 + math.cos(math.pi * c.lr_warm_epochs / c.meta_epochs)) / 2
- else:
- c.lr_warmup_to = c.lr
- context.set_context(mode=context.GRAPH_MODE, device_target=c.device_target)
- context.set_context(device_id=c.device_id)
- init_seeds(seed=int(time.time()))
-
-
- #export
- if(c.export_mode == 'encoder'):
- model, pool_layers, pool_dims = load_encoder_arch(c, c.L)
- param_dict = load_checkpoint(c.exported_ckpt_path+"encoder.ckpt")
- load_param_into_net(model, param_dict)
- input_arr = Tensor(np.ones([c.batch_size, 3, c.input_size, c.input_size]), ms.float32)#32,3,256,256
- export(model, input_arr, file_name=c.exported_file_name+"_encoder", file_format=c.file_format)
- else:
- #export 3 decoders
- model_1 = load_decoder_arch(c, 512)
- model_2 = load_decoder_arch(c, 1024)
- model_3 = load_decoder_arch(c, 2048)
- param_dict_1 = load_checkpoint(c.exported_ckpt_path+"decoder_0.ckpt")
- param_dict_2 = load_checkpoint(c.exported_ckpt_path+"decoder_1.ckpt")
- param_dict_3 = load_checkpoint(c.exported_ckpt_path+"decoder_2.ckpt")
- load_param_into_net(model_1, param_dict_1)
- load_param_into_net(model_2, param_dict_2)
- load_param_into_net(model_3, param_dict_3)
- epcp_shape_list = [[256,512], [256,128], [256,1024], [256,128], [64,2048], [64,128]] \
- if c.input_size == 256 else \
- [[256,512], [256,128], [256,1024], [256,128], [256,2048], [256,128]]
- #c.input_size must in [256,512]
- print('epcp shape list',epcp_shape_list)
- input_arr_1_1 = Tensor(np.ones(epcp_shape_list[1]), ms.float32)
- input_arr_1_0 = Tensor(np.ones(epcp_shape_list[0]), ms.float32)
- input_arr_2_1 = Tensor(np.ones(epcp_shape_list[3]), ms.float32)
- input_arr_2_0 = Tensor(np.ones(epcp_shape_list[2]), ms.float32)
- input_arr_3_1 = Tensor(np.ones(epcp_shape_list[5]), ms.float32)
- input_arr_3_0 = Tensor(np.ones(epcp_shape_list[4]), ms.float32)
-
- inputs_1 = [input_arr_1_0, input_arr_1_1]
- inputs_2 = [input_arr_2_0, input_arr_2_1]
- inputs_3 = [input_arr_3_0, input_arr_3_1]
- export(model_1, *inputs_1, file_name=c.exported_file_name + "_decoder_0",file_format=c.file_format)
- export(model_2, *inputs_2, file_name=c.exported_file_name + "_decoder_1",file_format=c.file_format)
- export(model_3, *inputs_3, file_name=c.exported_file_name + "_decoder_2",file_format=c.file_format)
|