When I train high_resolution image by efficientAD with input size（1024*1024），it host too much GPU memory ,So I modified the torch_model.py in EfficientAD.

[StarShang]

When I train high_resolution image like 50004000 by efficientAD with input size（10241024），the memory on GPU host too much.

this is the config file：

dataset:
  name: mvtec
  format: mvtec
  path: ./datasets/MVTec
  category: bigimage
  task: segmentation
  train_batch_size: 1
  eval_batch_size: 16
  num_workers: 8
  image_size: [1024,1024] # dimensions to which images are resized (mandatory)
  center_crop: null # dimensions to which images are center-cropped after resizing (optional)
  normalization: none # data distribution to which the images will be normalized: [none, imagenet]
  transform_config:
    train: null
    eval: null
  test_split_mode: from_dir # options: [from_dir, synthetic]
  test_split_ratio: 0.2 # fraction of train images held out testing (usage depends on test_split_mode)
  val_split_mode: same_as_test # options: [same_as_test, from_test, synthetic]
  val_split_ratio: 0.5 # fraction of train/test images held out for validation (usage depends on val_split_mode)

model:
  name: efficient_ad
  teacher_out_channels: 384
  model_size: small # options: [small, medium]
  lr: 0.0001
  weight_decay: 0.00001
  padding: false
  pad_maps: true # relevant for "padding: false", see EfficientAd in lightning_model.py
  # generic params
  normalization_method: min_max # options: [null, min_max, cdf]

metrics:
  image:
    - AUROC
    - F1Score
  pixel:
    - AUROC
    - AUPRO
  threshold:
    method: adaptive #options: [adaptive, manual]
    manual_image: null
    manual_pixel: null

visualization:
  show_images: False # show images on the screen
  save_images: True # save images to the file system
  log_images: False # log images to the available loggers (if any)
  image_save_path: null # path to which images will be saved
  mode: full # options: ["full", "simple"]

project:
  seed: 42
  path: ./results

logging:
  logger: [] # options: [comet, tensorboard, wandb, csv] or combinations.
  log_graph: false # Logs the model graph to respective logger.

optimization:
  export_mode: null # options: torch, onnx, openvino
# PL Trainer Args. Don't add extra parameter here.
trainer:
  enable_checkpointing: true
  default_root_dir: null
  gradient_clip_val: 0
  gradient_clip_algorithm: norm
  num_nodes: 1
  devices: 1
  enable_progress_bar: true
  overfit_batches: 0.0
  track_grad_norm: -1
  check_val_every_n_epoch: 1
  fast_dev_run: false
  accumulate_grad_batches: 1
  max_epochs: 200
  min_epochs: null
  max_steps: 70000
  min_steps: null
  max_time: null
  limit_train_batches: 1.0
  limit_val_batches: 1.0
  limit_test_batches: 1.0
  limit_predict_batches: 1.0
  val_check_interval: 1.0
  log_every_n_steps: 50
  accelerator: auto # <"cpu", "gpu", "tpu", "ipu", "hpu", "auto">
  strategy: dp
  sync_batchnorm: false
  precision: 32
  enable_model_summary: true
  num_sanity_val_steps: 0
  profiler: null
  benchmark: false
  deterministic: false
  reload_dataloaders_every_n_epochs: 0
  auto_lr_find: false
  replace_sampler_ddp: true
  detect_anomaly: false
  auto_scale_batch_size: false
  plugins: null
  move_metrics_to_cpu: false
  multiple_trainloader_mode: max_size_cycle

So I changed the torch_model.py in EfficientAD like this:

_def forward(self, batch: Tensor, batch_imagenet: Tensor = None) -> Tensor | dict:
       """Prediction by EfficientAd models.

       Args:
           batch (Tensor): Input images.

       Returns:
           Tensor: Predictions
       """
       with torch.no_grad():
           teacher_output = self.teacher(batch)
           if self.is_set(self.mean_std):
               teacher_output = (teacher_output - self.mean_std["mean"]) / self.mean_std["std"]
----   student_output = self.student(batch)
++++   with torch.no_grad():
++++        student_output = self.student(batch)
       distance_st = torch.pow(teacher_output - student_output[:, : self.teacher_out_channels, :, :], 2)

       if self.training:
           # Student loss
           distance_st = reduce_tensor_elems(distance_st)
           d_hard = torch.quantile(distance_st, 0.999)
           loss_hard = torch.mean(distance_st[distance_st >= d_hard])
----       student_output_penalty = self.student(batch_imagenet)[:, : self.teacher_out_channels, :, :]
++++       with torch.no_grad():
++++             student_output_penalty = self.student(batch_imagenet)[:, : self.teacher_out_channels, :, :]
           loss_penalty = torch.mean(student_output_penalty**2)
           loss_st = loss_hard + loss_penalty

           # Autoencoder and Student AE Loss
           aug_img = self.choose_random_aug_image(batch)
           ae_output_aug = self.ae(aug_img)

           with torch.no_grad():
               teacher_output_aug = self.teacher(aug_img)
               if self.is_set(self.mean_std):
                   teacher_output_aug = (teacher_output_aug - self.mean_std["mean"]) / self.mean_std["std"]

           student_output_ae_aug = self.student(aug_img)[:, self.teacher_out_channels :, :, :]

           distance_ae = torch.pow(teacher_output_aug - ae_output_aug, 2)
           distance_stae = torch.pow(ae_output_aug - student_output_ae_aug, 2)

           loss_ae = torch.mean(distance_ae)
           loss_stae = torch.mean(distance_stae)
           return (loss_st, loss_ae, loss_stae)_

After I changed the code which is used for self.student to forward ,the GPU memory host reduced by 20%,but the pixel_AUROC reduced from 0.98 to 0.93 Did I make a mistake!？ Also， May I ask if there is a good way to perform anomaly detection on high-resolution images？

[NilsB98]

Hi @StarShang,

since I'm just starting out with the library as well I'm not quite sure whether this is applicaple for EfficientAD, but there is a class called 'Tiler' in the data.utils module. This will split your large image into multiple tiles. Maybe this also helps in this case?

Regarding the implementation that you gave: I think the students weights should be updated and therefore the gradient calculated, so I'm not quite sure why you run it with torch.no_grad().
Would you mind giving some insight about your reasoning for this change?

[blaz-r]

You are right. torch.no_grad() can't be used there as that will prevent network from learning.

Regarding the processing of large resolution images, you could use the tiler, but the default tiler does not in fact reduce the memory usage. We are working on the tiled ensemble that will soon be published that will enable efficient processing of high resolution images. But right now, you can reduce the size or you can split the images into tiles in advance.

[StarShang]

Thanks @NilsB98 @blaz-r I have revised the model again, and in each forward propagation of three self.students, each epoch only allows gradient descent along one self.student. I trained and found that the memory usage decreased by about 20%, and the final accuracy decreased from 98% to 97.7%.

[blaz-r]

The other student pass is crucial to the process so it doesn't overfit to data at hand. Looks like the accuracy is still good, so I think it's okay, but as I said keep in mind the second pass is to prevent overfitting.