From 4bf3d50e54ff86005c0e8b59db72405c27cee3f1 Mon Sep 17 00:00:00 2001 From: Tekno-H <95920362+Tekno-H@users.noreply.github.com> Date: Tue, 13 Dec 2022 09:26:36 +0200 Subject: [PATCH 1/2] Create cflow-to-onnx.py Script to export the Cflow-ad model to ONNX format. --- cflow-to-onnx.py | 196 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 196 insertions(+) create mode 100644 cflow-to-onnx.py diff --git a/cflow-to-onnx.py b/cflow-to-onnx.py new file mode 100644 index 0000000..3ce3fc3 --- /dev/null +++ b/cflow-to-onnx.py @@ -0,0 +1,196 @@ +import os, time +import random, math +import numpy as np +import torch +import torch.nn.functional as F +from sklearn.metrics import roc_auc_score, auc, precision_recall_curve +from skimage.measure import label, regionprops +from tqdm import tqdm +from visualize import * +from model import load_decoder_arch, load_encoder_arch, positionalencoding2d, activation +from utils import * +from custom_datasets import * +from custom_models import * +from config import get_args +from config import get_args +from train import train +from test import test +from test_single import test_single + +log_theta = torch.nn.LogSigmoid() +# First we create the model ex. model=CFlow(C) +# Then we load the parameters ex. model(checkpoint) + +def load_weights2(model, filename): + path = os.path.join(filename) + state = torch.load(path) + model.encoder.load_state_dict(state['encoder_state_dict'], strict=False) + decoders = [decoder.load_state_dict(state, strict=False) for decoder, state in zip(model.decoders, state['decoder_state_dict'])] + print('Loading weights from {}'.format(filename)) + return model + +class CFlow(torch.nn.Module): + def __init__(self, c): + super(CFlow, self).__init__() + L = c.pool_layers + self.encoder, self.pool_layers, self.pool_dims = load_encoder_arch(c, L) + self.encoder = self.encoder.to(c.device).eval() + self.decoders = [load_decoder_arch(c, pool_dim) for pool_dim in self.pool_dims] + self.decoders = [decoder.to(c.device) for decoder in self.decoders] + params = list(self.decoders[0].parameters()) + for l in range(1, L): + params += list(self.decoders[l].parameters()) + # optimizer + self.optimizer = torch.optim.Adam(params, lr=c.lr) + self.N=256 + + def forward(self, x): + print(type(x)) + P = c.condition_vec + #print(self.decoders) + self.decoders = [decoder.eval() for decoder in self.decoders] + height = list() + width = list() + i=0 + test_dist = [list() for layer in self.pool_layers] + test_loss = 0.0 + test_count = 0 + start = time.time() + _ = self.encoder(x) + with torch.no_grad(): + for l, layer in enumerate(self.pool_layers): + e = activation[layer] # BxCxHxW + # + B, C, H, W = e.size() + S = H * W + E = B * S + # + if i == 0: # get stats + height.append(H) + width.append(W) + # + p = positionalencoding2d(P, H, W).to(c.device).unsqueeze(0).repeat(B, 1, 1, 1) + c_r = p.reshape(B, P, S).transpose(1, 2).reshape(E, P) # BHWxP + e_r = e.reshape(B, C, S).transpose(1, 2).reshape(E, C) # BHWxC + + decoder = self.decoders[l] + FIB = E // self.N + int(E % self.N > 0) # number of fiber batches + for f in range(FIB): + if f < (FIB - 1): + idx = torch.arange(f * self.N, (f + 1) * self.N) + else: + idx = torch.arange(f * self.N, E) + # + c_p = c_r[idx] # NxP + e_p = e_r[idx] # NxC + # m_p = m_r[idx] > 0.5 # Nx1 + # + if 'cflow' in c.dec_arch: + z, log_jac_det = decoder(e_p, [c_p, ]) + else: + z, log_jac_det = decoder(e_p) + # + decoder_log_prob = get_logp(C, z, log_jac_det) + log_prob = decoder_log_prob / C # likelihood per dim + loss = -log_theta(log_prob) + test_loss += t2np(loss.sum()) + test_count += len(loss) + test_dist[l] = test_dist[l] + log_prob.detach().cpu().tolist() + return height, width, test_dist + +def init_seeds(seed=0): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def main(c): + # model + c.gpu='0' + c.enc_arch='mobilenet_v3_large' + c.inp=256 + c.dataset='mvtec' + c.action_type='norm-train' + + # image + c.img_size = (c.input_size, c.input_size) # HxW format + c.crp_size = (c.input_size, c.input_size) # HxW format + if c.dataset == 'stc': + c.norm_mean, c.norm_std = 3 * [0.5], 3 * [0.225] + else: + c.norm_mean, c.norm_std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] + # + c.img_dims = [3] + list(c.img_size) + # network hyperparameters + c.clamp_alpha = 1.9 # see paper equation 2 for explanation + c.condition_vec = 128 + c.dropout = 0.0 # dropout in s-t-networks + # dataloader parameters + if c.dataset == 'mvtec': + c.data_path = './data/MVTec-AD' + elif c.dataset == 'stc': + c.data_path = './data/STC/shanghaitech' + elif c.dataset == 'video': + c.data_path = c.video_path + elif c.dataset == 'image': + c.data_path = c.image_path + + else: + raise NotImplementedError('{} is not supported dataset!'.format(c.dataset)) + # output settings + 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]] + print('LR schedule: {}'.format(c.lr_decay_epochs)) + 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 + ######## + os.environ['CUDA_VISIBLE_DEVICES'] = c.gpu + c.use_cuda = not c.no_cuda and torch.cuda.is_available() + init_seeds(seed=int(time.time())) + c.device = torch.device("cuda" if c.use_cuda else "cpu") + # selected function: + + model = CFlow(c) + print("Created !!!") + PATH = 'weights/mvtec_mobilenet_v3_large_freia-cflow_pl3_cb8_inp256_run0_Model_2022-11-08-10:50:39.pt' + model=load_weights2(model,PATH) + print("Loaded !") + model.eval() + + batch_size = 1 + x = torch.randn(batch_size, 3, 256, 256).to(c.device) + out=model(x) + print("Starting Export !!!") + torch.onnx.export( + model, + x, + "custom.onnx", + export_params=True, + verbose=True, + opset_version=11, + input_names=["input"], + output_names=["output"], + ) + + + +if __name__ == '__main__': + c = get_args() + main(c) From 37dc8d305af3361522ea878a39d033ed8145d92f Mon Sep 17 00:00:00 2001 From: Tekno-H <95920362+Tekno-H@users.noreply.github.com> Date: Tue, 13 Dec 2022 09:31:17 +0200 Subject: [PATCH 2/2] Create cflow-to-onnx2.py Another script to export the model ONNX format that gives the same error as the previous one --- cflow-to-onnx2.py | 212 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 212 insertions(+) create mode 100644 cflow-to-onnx2.py diff --git a/cflow-to-onnx2.py b/cflow-to-onnx2.py new file mode 100644 index 0000000..c1f48ba --- /dev/null +++ b/cflow-to-onnx2.py @@ -0,0 +1,212 @@ +import os, time +import random, math +import numpy as np +import torch +import torch.nn.functional as F +from sklearn.metrics import roc_auc_score, auc, precision_recall_curve +from skimage.measure import label, regionprops +from tqdm import tqdm +from visualize import * +from model import load_decoder_arch, load_encoder_arch, positionalencoding2d, activation +from utils import * +from custom_datasets import * +from custom_models import * +from config import get_args +from config import get_args +from train import train +from test import test +from test_single import test_single + +log_theta = torch.nn.LogSigmoid() + +def load_weights2(model, filename): + path = os.path.join(filename) + state = torch.load(path) + model.Encoder_module.encoder.load_state_dict(state['encoder_state_dict'], strict=False) + decoders = [decoder.load_state_dict(state, strict=False) for decoder, state in zip(model.Decoder_module.decoders, state['decoder_state_dict'])] + print('Loading weights from {}'.format(filename)) + return model + +# Define the encoder +class Encoder(torch.nn.Module): + def __init__(self, encoder): + super(Encoder,self).__init__() + self.encoder = encoder + def forward(self, input): + return self.encoder(input) + +#Define the decoder +class Decoder(torch.nn.Module): + def __init__(self, c, decoders): + super(Decoder, self).__init__() + self.c = c + self.decoders = decoders + L = c.pool_layers + params = list(self.decoders[0].parameters()) + for l in range(1, L): + params += list(self.decoders[l].parameters()) + # optimizer + self.optimizer = torch.optim.Adam(params, lr=self.c.lr) + self.N = 256 + + def forward(self, pool_layers): + P = self.c.condition_vec + # print(self.decoders) + self.decoders = [decoder.eval() for decoder in self.decoders] + height = list() + width = list() + i = 0 + test_dist = [list() for layer in pool_layers] + test_loss = 0.0 + test_count = 0 + start = time.time() + with torch.no_grad(): + for l, layer in enumerate(pool_layers): + e = activation[layer] # BxCxHxW + # + B, C, H, W = e.size() + S = H * W + E = B * S + # + if i == 0: # get stats + height.append(H) + width.append(W) + # + p = positionalencoding2d(P, H, W).to(self.c.device).unsqueeze(0).repeat(B, 1, 1, 1) + c_r = p.reshape(B, P, S).transpose(1, 2).reshape(E, P) # BHWxP + e_r = e.reshape(B, C, S).transpose(1, 2).reshape(E, C) # BHWxC + + decoder = self.decoders[l] + FIB = E // self.N + int(E % self.N > 0) # number of fiber batches + for f in range(FIB): + if f < (FIB - 1): + idx = torch.arange(f * self.N, (f + 1) * self.N) + else: + idx = torch.arange(f * self.N, E) + # + c_p = c_r[idx] # NxP + e_p = e_r[idx] # NxC + # m_p = m_r[idx] > 0.5 # Nx1 + # + z, log_jac_det = decoder(e_p, [c_p, ]) + # + decoder_log_prob = get_logp(C, z, log_jac_det) + log_prob = decoder_log_prob / C # likelihood per dim + loss = -log_theta(log_prob) + test_loss += t2np(loss.sum()) + test_count += len(loss) + test_dist[l] = test_dist[l] + log_prob.detach().cpu().tolist() + return height, width, test_dist + +#Define the base calss for the model +class CFlow(torch.nn.Module): + def __init__(self, c,encoder,decoders,pool_layers): + super(CFlow, self).__init__() + self.pool_layers=pool_layers + self.Encoder_module = Encoder(encoder) + self.Decoder_module = Decoder(c, decoders) + def forward(self,enc_input): + _=self.Encoder_module(enc_input) + height, width, test_dist = self.Decoder_module(self.pool_layers) + return height, width , test_dist + + +def init_seeds(seed=0): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) +#Main +def main(c): + # model + c.gpu = '0' + c.enc_arch = 'mobilenet_v3_large' + c.inp = 256 + c.dataset = 'mvtec' + c.action_type = 'norm-train' + + # image + c.img_size = (c.input_size, c.input_size) # HxW format + c.crp_size = (c.input_size, c.input_size) # HxW format + if c.dataset == 'stc': + c.norm_mean, c.norm_std = 3 * [0.5], 3 * [0.225] + else: + c.norm_mean, c.norm_std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] + # + c.img_dims = [3] + list(c.img_size) + # network hyperparameters + c.clamp_alpha = 1.9 # see paper equation 2 for explanation + c.condition_vec = 128 + c.dropout = 0.0 # dropout in s-t-networks + # dataloader parameters + if c.dataset == 'mvtec': + c.data_path = './data/MVTec-AD' + elif c.dataset == 'stc': + c.data_path = './data/STC/shanghaitech' + elif c.dataset == 'video': + c.data_path = c.video_path + elif c.dataset == 'image': + c.data_path = c.image_path + else: + raise NotImplementedError('{} is not supported dataset!'.format(c.dataset)) + # output settings + 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]] + print('LR schedule: {}'.format(c.lr_decay_epochs)) + 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 + ######## + os.environ['CUDA_VISIBLE_DEVICES'] = c.gpu + c.use_cuda = not c.no_cuda and torch.cuda.is_available() + init_seeds(seed=int(time.time())) + c.device = torch.device("cuda" if c.use_cuda else "cpu") + #Create the encoder and decoder networks + L = c.pool_layers + encoder, pool_layers, pool_dims = load_encoder_arch(c, L) + encoder = encoder.to(c.device).eval() + decoders = [load_decoder_arch(c, pool_dim) for pool_dim in pool_dims] + decoders = [decoder.to(c.device) for decoder in decoders] + + #Initialize the base calss + model=CFlow(c,encoder,decoders,pool_layers) + + PATH = 'weights/mvtec_mobilenet_v3_large_freia-cflow_pl3_cb8_inp256_run0_Model_2022-11-08-10:50:39.pt' + model = load_weights2(model, PATH) + print("Loaded !") + model.eval() + batch_size = 1 + x = torch.randn(batch_size, 3, 256, 256).to(c.device) + out = model(x) + torch.onnx.export( + model, # + x, + "custom-d.onnx", + export_params=True, + verbose=True, + opset_version=11, + input_names=["input"], + output_names=["output"], + ) +if __name__ == '__main__': + c = get_args() + main(c) + + + +