devbox
花生_TL007
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

自己实现SimpleNet_simplenet代码复现

作者：花生_TL007 | 2024-05-05 07:47:22

赞

踩

simplenet代码复现

数据加载：

  1. import os
  2. import PIL
  3. import torch
  4. from torchvision import transforms
  5.  
  6. IMAGENET_MEAN = [0.485, 0.456, 0.406]
  7. IMAGENET_STD = [0.229, 0.224, 0.225]
  8.  
  9. class Dataset(torch.utils.data.Dataset):
  10.     def __init__(
  11.         self,
  12.         source,
  13.         resize=512,
  14.         imagesize=512,
  15.         rotate_degrees=0,
  16.         translate=0,
  17.         brightness_factor=0,
  18.         contrast_factor=0,
  19.         saturation_factor=0,
  20.         gray_p=0,
  21.         h_flip_p=0,
  22.         v_flip_p=0,
  23.         scale=0,
  24.         **kwargs,
  25.     ):
  26.         super().__init__()
  27.         self.source = source
  28.         self.imgpathlist = [os.path.join(self.source,filename) for filename in os.listdir(self.source)]
  29.         self.transform_img = [
  30.             transforms.Resize(resize),
  31.             # transforms.RandomRotation(rotate_degrees, transforms.InterpolationMode.BILINEAR),
  32.             transforms.ColorJitter(brightness_factor, contrast_factor, saturation_factor),
  33.             transforms.RandomHorizontalFlip(h_flip_p),
  34.             transforms.RandomVerticalFlip(v_flip_p),
  35.             transforms.RandomGrayscale(gray_p),
  36.             transforms.RandomAffine(rotate_degrees,
  37.                                     translate=(translate, translate),
  38.                                     scale=(1.0-scale, 1.0+scale),
  39.                                     interpolation=transforms.InterpolationMode.BILINEAR),
  40.             transforms.CenterCrop(imagesize),
  41.             transforms.ToTensor(),
  42.             transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
  43.         ]
  44.         self.transform_img = transforms.Compose(self.transform_img)
  45.  
  46.         self.transform_mask = [
  47.             transforms.Resize(resize),
  48.             transforms.CenterCrop(imagesize),
  49.             transforms.ToTensor(),
  50.         ]
  51.         self.transform_mask = transforms.Compose(self.transform_mask)
  52.         self.imagesize = (3, imagesize, imagesize)
  53.  
  54.     def __getitem__(self, idx):
  55.         image = PIL.Image.open(self.imgpathlist[idx]).convert("RGB")
  56.         image = self.transform_img(image)
  57.         mask = torch.zeros([1, *image.size()[1:]])
  58.         return image,mask
  59.  
  60.     def __len__(self):
  61.         return len(self.imgpathlist)
  62.  
  63. # dataset = Dataset(source=r"D:\dataset\mvtec_anomaly_detection\bottle\train\good")
  64. # image,mask = dataset.__getitem__(0)
  65. # print(image.shape,mask.shape)

网络：

  1. import torch
  2. import math
  3. from torch import onnx
  4. from torch import nn
  5. import numpy as np
  6. import torch.nn.functional as F
  7. import scipy.ndimage as ndimage
  8. import backbones
  9. import copy
  10.  
  11. def pth2onnx(model, dummy_input, dynamiconnx):
  12.     torch.set_grad_enabled(False)
  13.     input_names = ["input1"]
  14.     output_names = ["output1"]
  15.     # 保存维度变化的onnx
  16.     onnx.export(model=model, args=dummy_input, f=dynamiconnx, input_names=input_names,
  17.                 output_names=output_names, verbose=False,
  18.                 dynamic_axes=dict([(k, {0: 'batch_size'}) for k in input_names] +
  19.                                   [(k, {0: 'batch_size'}) for k in output_names]),
  20.                 keep_initializers_as_inputs=True)
  21.  
  22. class Convolution(torch.nn.Module):
  23.     def __init__(self,in_chanel,out_chanel,kernalsize,strid,padding):
  24.         super(Convolution,self).__init__()
  25.         self.conv = nn.Conv2d(in_chanel,out_chanel,kernalsize,strid,padding)
  26.         self.bn = nn.BatchNorm2d(out_chanel)
  27.         self.active = nn.Mish(True)
  28.  
  29.     def forward(self,x):
  30.         return self.active(self.bn(self.conv(x)))
  31.  
  32. class PatchMaker(nn.Module):
  33.     def __init__(self, patchsize, top_k=0, stride=None):
  34.         super(PatchMaker,self).__init__()
  35.         self.patchsize = patchsize
  36.         self.stride = stride
  37.         self.top_k = top_k
  38.  
  39.     def forward(self, features):
  40.         """Convert a tensor into a tensor of respective patches.
  41.         Args:
  42.             x: [torch.Tensor, bs x c x w x h]
  43.         Returns:
  44.             x: [torch.Tensor, bs * w//stride * h//stride, c, patchsize,
  45.             patchsize]
  46.         """
  47.         return_spatial_info = True
  48.         padding = int((self.patchsize - 1) / 2)#1
  49.         unfolder = torch.nn.Unfold(
  50.             kernel_size=self.patchsize, stride=self.stride, padding=padding, dilation=1
  51.         )
  52.         unfolded_features = unfolder(features)
  53.         number_of_total_patches = []
  54.         for s in features.shape[-2:]:
  55.             n_patches = (
  56.                 s + 2 * padding - 1 * (self.patchsize - 1) - 1
  57.             ) / self.stride + 1
  58.             number_of_total_patches.append(int(n_patches))
  59.         unfolded_features = unfolded_features.reshape(
  60.             *features.shape[:2], self.patchsize, self.patchsize, -1
  61.         )
  62.         unfolded_features = unfolded_features.permute(0, 4, 1, 2, 3)
  63.  
  64.         if return_spatial_info:
  65.             return unfolded_features, number_of_total_patches
  66.         return unfolded_features
  67.  
  68. class Resblock(nn.Module):
  69.     def __init__(self,ch):
  70.         super(Resblock,self).__init__()
  71.         self.conv1 = Convolution(ch, ch // 2, 1, 1, 0)
  72.         self.conv2 = Convolution(ch // 2,ch // 2, 3, 1, 1)
  73.         self.conv3 = nn.Conv2d(ch // 2,ch, 1, 1)
  74.         self.relu = nn.ReLU(True)
  75.     def forward(self,x):
  76.         y = self.conv1(x)
  77.         y = self.conv2(y)
  78.         y = self.conv3(y)
  79.         return self.relu(x + y)
  80.  
  81. class Preprocessing(torch.nn.Module):
  82.     def __init__(self, input_dims, output_dim):
  83.         super(Preprocessing, self).__init__()
  84.         self.input_dims = input_dims
  85.         self.output_dim = output_dim
  86.         self.preprocessing_modules = torch.nn.ModuleList()
  87.         for input_dim in input_dims:
  88.             module = MeanMapper(output_dim)
  89.             self.preprocessing_modules.append(module)
  90.  
  91.     def forward(self, features):
  92.         _features = []
  93.         for module, feature in zip(self.preprocessing_modules, features):
  94.             _features.append(module(feature))
  95.         return torch.stack(_features, dim=1)
  96.  
  97. class MeanMapper(torch.nn.Module):
  98.     def __init__(self, preprocessing_dim):
  99.         super(MeanMapper, self).__init__()
  100.         self.preprocessing_dim = preprocessing_dim
  101.     def forward(self, features):
  102.         features = features.reshape(len(features), 1, -1)
  103.         return F.adaptive_avg_pool1d(features, self.preprocessing_dim).squeeze(1)
  104.  
  105. def init_weight(m):
  106.  
  107.     if isinstance(m, torch.nn.Linear):
  108.         torch.nn.init.xavier_normal_(m.weight)
  109.     elif isinstance(m, torch.nn.Conv2d):
  110.         torch.nn.init.xavier_normal_(m.weight)
  111.  
  112. class Projection(torch.nn.Module):
  113.     def __init__(self, in_planes, out_planes=None, n_layers=1, layer_type=0):
  114.         super(Projection, self).__init__()
  115.  
  116.         if out_planes is None:
  117.             out_planes = in_planes
  118.         self.layers = torch.nn.Sequential()
  119.         _in = None
  120.         _out = None
  121.         for i in range(n_layers):
  122.             _in = in_planes if i == 0 else _out
  123.             _out = out_planes
  124.             self.layers.add_module(f"{i}fc",
  125.                                    torch.nn.Linear(_in, _out))
  126.             if i < n_layers - 1:
  127.                 # if layer_type > 0:
  128.                 #     self.layers.add_module(f"{i}bn",
  129.                 #                            torch.nn.BatchNorm1d(_out))
  130.                 if layer_type > 1:
  131.                     self.layers.add_module(f"{i}relu",
  132.                                            torch.nn.LeakyReLU(.2))
  133.         self.apply(init_weight)
  134.  
  135.     def forward(self, x):
  136.         # x = .1 * self.layers(x) + x
  137.         x = self.layers(x)
  138.         return x
  139.  
  140. class Discriminator(torch.nn.Module):
  141.     def __init__(self, in_planes, n_layers=1, hidden=None):
  142.         super(Discriminator, self).__init__()
  143.         _hidden = in_planes if hidden is None else hidden
  144.         self.body = torch.nn.Sequential()
  145.         for i in range(n_layers-1):
  146.             _in = in_planes if i == 0 else _hidden
  147.             _hidden = int(_hidden // 1.5) if hidden is None else hidden
  148.             self.body.add_module('block%d'%(i+1),
  149.                                  torch.nn.Sequential(
  150.                                      torch.nn.Linear(_in, _hidden),
  151.                                      torch.nn.BatchNorm1d(_hidden),
  152.                                      torch.nn.LeakyReLU(0.2)
  153.                                  ))
  154.         self.tail = torch.nn.Linear(_hidden, 1, bias=False)
  155.         self.apply(init_weight)
  156.  
  157.     def forward(self,x):
  158.         x = self.body(x)
  159.         x = self.tail(x)
  160.         return x
  161.  
  162. class ForwardHook:
  163.     def __init__(self, hook_dict, layer_name: str, last_layer_to_extract: str):
  164.         self.hook_dict = hook_dict
  165.         self.layer_name = layer_name
  166.         self.raise_exception_to_break = copy.deepcopy(
  167.             layer_name == last_layer_to_extract
  168.         )
  169.  
  170.     def __call__(self, module, input, output):
  171.         self.hook_dict[self.layer_name] = output
  172.         return None
  173.  
  174. class NetworkFeatureAggregator(torch.nn.Module):
  175.     """Efficient extraction of network features."""
  176.  
  177.     def __init__(self, backbone, layers_to_extract_from, device, train_backbone=False):
  178.         super(NetworkFeatureAggregator, self).__init__()
  179.         """Extraction of network features.
  181.         Runs a network only to the last layer of the list of layers where
  182.         network features should be extracted from.
  184.         Args:
  185.             backbone: torchvision.model
  186.             layers_to_extract_from: [list of str]
  187.         """
  188.         self.layers_to_extract_from = layers_to_extract_from
  189.         self.backbone = backbone
  190.         self.device = device
  191.         self.train_backbone = train_backbone
  192.  
  193.         for extract_layer in layers_to_extract_from:
  194.             if extract_layer == 'layer2':
  195.                 self.network_layer2 = backbone.__dict__["_modules"][extract_layer]
  196.             if extract_layer == 'layer3':
  197.                 self.network_layer3 = backbone.__dict__["_modules"][extract_layer]
  198.         print(self.network_layer2,'#%$^&*^*&(^^$%\n',self.network_layer3)
  199.         self.to(self.device)
  200.  
  201.     def forward(self, images, eval=True):
  202.         y = torch.randn((1,1000))
  203.         if self.train_backbone and not eval:
  204.             y = self.backbone(images)
  205.         else:
  206.             with torch.no_grad():
  207.                 try:
  208.                     y = self.backbone(images)
  209.                 except:
  210.                     pass
  211.         return y
  212.  
  213. class unsupervisedNet(torch.nn.Module):
  214.     def __init__(self,batchsize,train):
  215.         super(unsupervisedNet,self).__init__()
  216.         self.batchsize = batchsize
  217.         self.train_backbone = train
  218.         self.backbone_name = 'wideresnet50'
  219.         self.backbone = backbones.load(self.backbone_name)
  220.         self.device = 'cpu'  # 0
  221.         self.input_shape = [3, 512, 512]
  222.  
  223.         self.target_size = self.input_shape[-2:]
  224.         self.patchsize = 3
  225.         self.stride = 1
  226.         self.top_k = 0
  227.         self.input_dims = [512, 1024]
  228.         self.output_dim = 1536
  229.         self.smoothing = 4
  230.  
  231.         # self.feature_dimensions = [512, 1024]
  232.         # self.preprocessing = Preprocessing(self.feature_dimensions,self.output_dim)
  233.  
  234.         self.preprocessing_modules = torch.nn.ModuleList()
  235.         for input_dim in self.input_dims:
  236.             module = MeanMapper(self.output_dim)
  237.             self.preprocessing_modules.append(module)
  238.         self.pre_projection = Projection(self.output_dim, self.output_dim, 1, 0)
  239.  
  240.         self.discriminator = Discriminator(self.output_dim, n_layers=2, hidden=1024)
  241.  
  242.         self.layer2 = nn.Sequential(
  243.             self.backbone.conv1,
  244.             self.backbone.bn1,
  245.             self.backbone.relu,
  246.             self.backbone.maxpool,
  247.             self.backbone.layer1,
  248.             self.backbone.layer2,
  249.         )
  250.  
  251.         self.layer3 = nn.Sequential(
  252.             self.backbone.conv1,
  253.             self.backbone.bn1,
  254.             self.backbone.relu,
  255.             self.backbone.maxpool,
  256.             self.backbone.layer1,
  257.             self.backbone.layer2,
  258.             self.backbone.layer3
  259.         )
  260.  
  261.         self.unfolded_features = []
  262.         self.patch_shapes = []
  263.         self.padding = int((self.patchsize - 1) / 2)
  264.         self.unfolder = nn.Unfold(kernel_size=self.patchsize, stride=self.stride, padding=self.padding, dilation=1)
  265.  
  266.     def score(self, x):
  267.         was_numpy = False
  268.         if isinstance(x, np.ndarray):
  269.             was_numpy = True
  270.             x = torch.from_numpy(x)
  271.         while x.ndim > 2:
  272.             x = torch.max(x, dim=-1).values
  273.         if x.ndim == 2:
  274.             if self.top_k > 1:
  275.                 x = torch.topk(x, self.top_k, dim=1).values.mean(1)
  276.             else:
  277.                 x = torch.max(x, dim=1).values
  278.         if was_numpy:
  279.             return x.numpy()
  280.         return x
  281.  
  282.     def forward(self,x):
  283.         output1 = self.layer2(x)
  284.         output2 = self.layer3(x)
  285.         unfolded_features1 = self.unfolder(output1)
  286.         patch_shapes1 = []
  287.         for s in output1.shape[-2:]:
  288.             n_patches = (
  289.                                 s + 2 * self.padding - 1 * (self.patchsize - 1) - 1
  290.                         ) / self.stride + 1
  291.             patch_shapes1.append(int(n_patches))
  292.         unfolded_features1 = unfolded_features1.reshape(
  293.             *output1.shape[:2], self.patchsize, self.patchsize, -1
  294.         )
  295.         unfolded_features1 = unfolded_features1.permute(0, 4, 1, 2, 3)
  296.  
  297.         unfolded_features2 = self.unfolder(output2)
  298.         patch_shapes2 = []
  299.         for s in output2.shape[-2:]:
  300.             n_patches = (
  301.                                 s + 2 * self.padding - 1 * (self.patchsize - 1) - 1
  302.                         ) / self.stride + 1
  303.             patch_shapes2.append(int(n_patches))
  304.         unfolded_features2 = unfolded_features2.reshape(
  305.             *output2.shape[:2], self.patchsize, self.patchsize, -1
  306.         )
  307.         unfolded_features2 = unfolded_features2.permute(0, 4, 1, 2, 3)
  308.  
  309.         ref_num_patches = patch_shapes1
  310.         _features = unfolded_features2
  311.         patch_dims = patch_shapes2
  312.         _features = _features.reshape(
  313.             _features.shape[0], patch_dims[0], patch_dims[1], *_features.shape[2:]
  314.         )
  315.         _features = _features.permute(0, -3, -2, -1, 1, 2)
  316.         perm_base_shape = _features.shape
  317.         _features = _features.reshape(-1, *_features.shape[-2:])
  318.         _features = F.interpolate(
  319.             _features.unsqueeze(1),
  320.             size=(ref_num_patches[0], ref_num_patches[1]),
  321.             mode="bilinear",
  322.             align_corners=False,
  323.         )
  324.         _features = _features.squeeze(1)
  325.         _features = _features.reshape(
  326.             *perm_base_shape[:-2], ref_num_patches[0], ref_num_patches[1]
  327.         )
  328.         _features = _features.permute(0, -2, -1, 1, 2, 3)
  329.         _features = _features.reshape(len(_features), -1, *_features.shape[-3:])
  330.         unfolded_features2 = _features
  331.  
  332.         unfolded_features1 = unfolded_features1.reshape(-1, *unfolded_features1.shape[-3:])
  333.         unfolded_features2 = unfolded_features2.reshape(-1, *unfolded_features2.shape[-3:])
  334.  
  335.         # _features = []
  336.         model1 = self.preprocessing_modules[0]
  337.         feature1 = model1(unfolded_features1)
  338.         model2 = self.preprocessing_modules[1]
  339.         feature2 = model2(unfolded_features2)
  340.         features = torch.stack([feature1, feature2], dim=1)
  341.  
  342.         features = features.reshape(len(features), 1, -1)
  343.         features = F.adaptive_avg_pool1d(features, self.output_dim)
  344.         features = features.reshape(len(features), -1)
  345.  
  346.         patch_shapes = []
  347.         patch_shapes.append(patch_shapes1)
  348.         patch_shapes.append(patch_shapes2)
  349.  
  350.         features = self.pre_projection(features)  # torch.Size([4096, 1536])
  351.         self.features = features
  352.  
  353.         patch_scores = image_scores = -self.discriminator(features)
  354.  
  355.         patch_scores = patch_scores.cpu().detach().numpy()
  356.         image_scores = image_scores.cpu().detach().numpy()
  357.  
  358.         image_scores = image_scores.reshape(self.batchsize, -1, *image_scores.shape[1:])
  359.         image_scores = image_scores.reshape(*image_scores.shape[:2], -1)
  360.  
  361.         image_scores = self.score(image_scores)
  362.         image_scores = image_scores.reshape(self.batchsize, -1, *image_scores.shape[1:])
  363.         scales = patch_shapes[0]
  364.         # patch_scores = patch_scores.reshape(1, scales[0], scales[1])
  365.         patch_scores = patch_scores.reshape(self.batchsize, scales[0], scales[1])
  366.         # features = features.reshape(1, scales[0], scales[1], -1)
  367.         features = features.reshape(self.batchsize, scales[0], scales[1], -1)
  368.  
  369.         with torch.no_grad():
  370.             if isinstance(patch_scores, np.ndarray):
  371.                 patch_scores = torch.from_numpy(patch_scores)
  372.             _scores = patch_scores.to(self.device)
  373.             _scores = _scores.unsqueeze(1)
  374.             _scores = F.interpolate(
  375.                 _scores, size=self.target_size, mode="bilinear", align_corners=False
  376.             )
  377.             _scores = _scores.squeeze(1)
  378.             patch_scores = _scores.cpu().numpy()
  379.  
  380.             if isinstance(features, np.ndarray):
  381.                 features = torch.from_numpy(features)
  382.             features = features.to(self.device).permute(0, 3, 1, 2)
  383.             if self.target_size[0] * self.target_size[1] * features.shape[0] * features.shape[1] >= 2 ** 31:
  384.                 subbatch_size = int((2 ** 31 - 1) / (self.target_size[0] * self.target_size[1] * features.shape[1]))
  385.                 interpolated_features = []
  386.                 for i_subbatch in range(int(features.shape[0] / subbatch_size + 1)):
  387.                     subfeatures = features[i_subbatch * subbatch_size:(i_subbatch + 1) * subbatch_size]
  388.                     subfeatures = subfeatures.unsuqeeze(0) if len(subfeatures.shape) == 3 else subfeatures
  389.                     subfeatures = F.interpolate(
  390.                         subfeatures, size=self.target_size, mode="bilinear", align_corners=False
  391.                     )
  392.                     interpolated_features.append(subfeatures)
  393.                 features = torch.cat(interpolated_features, 0)
  394.             else:
  395.                 features = F.interpolate(
  396.                     features, size=self.target_size, mode="bilinear", align_corners=False
  397.                 )
  398.             # features = features.cpu().detach().numpy()
  399.         masks = [ndimage.gaussian_filter(patch_score, sigma=self.smoothing) for patch_score in patch_scores]
  400.         masks = torch.tensor(masks)
  401.         return masks  # ,self.patch_shapes
  402.  
  403. # net = unsupervisedNet(1,False)
  404. # # # net.cuda()
  405. # x = torch.randn((1,3,512,512))#.cuda()
  406. # # y = net(x)
  407. # # print(y.shape)
  408. # # pth2onnx(net,x,'test.onnx')
  409. # trace_script_module = torch.jit.trace(net,x)
  410. # trace_script_module.save('net1.torchscript')

训练代码：

  1. from net import unsupervisedNet
  2. import torch
  3. from mvdataset import Dataset
  4. from torch.utils.data import DataLoader
  5.  
  6. if __name__ == '__main__':
  7.     lr = 1e-3
  8.     mix_noise = 1
  9.     noise_std = 0.015
  10.     dsc_margin = .5
  11.     x = torch.randn(2,3,512,512)
  12.     label = torch.zeros(x.shape[0],512,512)
  13.     model = unsupervisedNet(x.shape[0],True)
  14.     dataset = Dataset(source=r"D:\dataset\mvtec_anomaly_detection\bottle\train\good")
  15.     dataloader = DataLoader(dataset,batch_size=2,shuffle=True,num_workers=4)
  16.     model.discriminator.train()
  17.     model.pre_projection.train()
  18.     dsc_opt = torch.optim.Adam(model.discriminator.parameters(), lr=0.0002, weight_decay=1e-5)
  19.     proj_opt = torch.optim.AdamW(model.pre_projection.parameters(), lr * .1)
  20.     backbone_opt = torch.optim.AdamW(model.backbone.parameters(), lr)
  21.     # x = x.cuda()
  22.     # label = label.cuda()
  23.     # model.cuda()
  24.     for epoch in range(100):
  25.         for data,label in dataloader:
  26.             dsc_opt.zero_grad()
  27.             proj_opt.zero_grad()
  28.             mask = model(data)
  29.             true_feats = model.features
  30.             noise_idxs = torch.randint(0,mix_noise, torch.Size([true_feats.shape[0]]))
  31.             noise_one_hot = torch.nn.functional.one_hot(noise_idxs, num_classes=mix_noise)#.to(self.device)
  32.             noise = torch.stack([
  33.                 torch.normal(0, noise_std * 1.1 ** (k), true_feats.shape)
  34.                 for k in range(mix_noise)], dim=1)#.to(self.device)
  35.             noise = (noise * noise_one_hot.unsqueeze(-1)).sum(1)
  36.             fake_feats = true_feats + noise
  37.             scores = model.discriminator(torch.cat([true_feats, fake_feats]))
  38.             true_scores = scores[:len(true_feats)]
  39.             fake_scores = scores[len(fake_feats):]
  40.             th = dsc_margin
  41.             p_true = (true_scores.detach() >= th).sum() / len(true_scores)
  42.             p_fake = (fake_scores.detach() < -th).sum() / len(fake_scores)
  43.             true_loss = torch.clip(-true_scores + th, min=0)
  44.             fake_loss = torch.clip(fake_scores + th, min=0)
  45.             loss = true_loss.mean() + fake_loss.mean()
  46.             print(loss)
  47.             loss.backward()
  48.             backbone_opt.step()
  49.             proj_opt.step()
  50.             dsc_opt.step()

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/花生_TL007/article/detail/537933
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号