Improve lot of types, imports and names

face_swapper/src/models/discriminator.py

@@ -2,7 +2,7 @@ import configparser
 from typing import List
 
 import numpy
-import torch.nn as nn
+from torch import nn
 
 from face_swapper.src.types import VisionTensor
 

face_swapper/src/models/generator.py

@@ -1,7 +1,7 @@
 import configparser
 from typing import Tuple
 
-import torch.nn as nn
+from torch import nn
 
 from face_swapper.src.networks.attribute_modulator import AADGenerator
 from face_swapper.src.networks.encoder import UNet

face_swapper/src/models/loss.py

@@ -18,7 +18,7 @@ class FaceSwapperLoss:
 		landmarker_path = CONFIG.get('training.model', 'landmarker_path')
 		motion_extractor_path = CONFIG.get('training.model', 'motion_extractor_path')
 		self.batch_size = CONFIG.getint('training.loader', 'batch_size')
-		self.mse_loss = torch.nn.MSELoss()
+		self.mse_loss = nn.MSELoss()
 		self.id_embedder = torch.jit.load(id_embedder_path, map_location = 'cpu')  # type:ignore[no-untyped-call]
 		self.landmarker = torch.jit.load(landmarker_path, map_location = 'cpu')  # type:ignore[no-untyped-call]
 		self.motion_extractor = torch.jit.load(motion_extractor_path, map_location = 'cpu')  # type:ignore[no-untyped-call]
@@ -127,7 +127,7 @@ class FaceSwapperLoss:
 
 	def get_face_landmarks(self, vision_tensor : VisionTensor) -> FaceLandmark203:
 		vision_tensor_norm = (vision_tensor + 1) * 0.5
-		vision_tensor_norm = torch.nn.functional.interpolate(vision_tensor_norm, size = (224, 224), mode = 'bilinear')
+		vision_tensor_norm = nn.functional.interpolate(vision_tensor_norm, size = (224, 224), mode = 'bilinear')
 		landmarks = self.landmarker(vision_tensor_norm)[2].view(-1, 203, 2)
 		return landmarks