class ImageConstant_Torch(input.ImageConstant): def __init__(self, image_path: str, imsize: int = 512, name: str = "ImageConstant_Torch", device=torch.device("cpu")): super(ImageConstant_Torch, self).__init__(name) loader = transforms.Compose( [transforms.Resize(imsize), transforms.ToTensor()]) image = loader(Image.open(image_path, mode="r")).unsqueeze(0) self.linked_tensor_torch = Tensor_Torch(image.to(device, torch.float), name=self.name + "_image_tensor") def get_saved_tensor(self): return self.linked_tensor_torch def set_device(self, device: torch.device): self.linked_tensor_torch.set_device(device=device) def get_device(self): return self.linked_tensor_torch.get_device() # return KB in memory usage for the loaded tensor def get_tensor_memory_size(self): return self.linked_tensor_torch.get_self_memory_size() # return KB in memory usage for gradients of the loaded tensor def get_tensor_grad_memory_size(self): return self.linked_tensor_torch.get_grad_memory_size() def remove_from_tracking_gradient(self): return self.linked_tensor_torch.remove_from_tracking_gradient() def start_tracking_gradient(self): return self.linked_tensor_torch.start_tracking_gradient() @staticmethod def get_description(): return "Loader for single image"
class ConstantConstant_Torch(input.ConstantConstant): def __init__(self, view: list, value: int, name: str = "ConstantConstant_Torch", device=torch.device("cpu")): super(ConstantConstant_Torch, self).__init__(name) self.linked_tensor_torch = Tensor_Torch( torch.add(torch.zeros(*view, device=device), value), name=self.name + "const_tensor") def get_saved_tensor(self): return self.linked_tensor_torch def set_device(self, device: torch.device): self.linked_tensor_torch.set_device(device=device) def get_device(self): return self.linked_tensor_torch.get_device() # return KB in memory usage for the loaded tensor def get_tensor_memory_size(self): return self.linked_tensor_torch.get_self_memory_size() # return KB in memory usage for gradients of the loaded tensor def get_tensor_grad_memory_size(self): return self.linked_tensor_torch.get_grad_memory_size() def remove_from_tracking_gradient(self): return self.linked_tensor_torch.remove_from_tracking_gradient() def start_tracking_gradient(self): return self.linked_tensor_torch.start_tracking_gradient() @staticmethod def get_description(): return "Constant tensor constant (1, 0)"
class TensorConstant_Torch(input.TensorConstant): def __init__(self, tensor_path: str, name: str = "TensorConstant_Torch", device=torch.device("cpu")): super(TensorConstant_Torch, self).__init__(name) self.linked_tensor_torch = Tensor_Torch( torch.load(tensor_path).to(device), name=self.name + "_saved_tensor") def get_saved_tensor(self): return self.linked_tensor_torch def set_device(self, device: torch.device): self.linked_tensor_torch.set_device(device=device) def get_device(self): return self.linked_tensor_torch.get_device() # return KB in memory usage for the loaded tensor def get_tensor_memory_size(self): return self.linked_tensor_torch.get_self_memory_size() # return KB in memory usage for gradients of the loaded tensor def get_tensor_grad_memory_size(self): return self.linked_tensor_torch.get_grad_memory_size() def remove_from_tracking_gradient(self): return self.linked_tensor_torch.remove_from_tracking_gradient() def start_tracking_gradient(self): return self.linked_tensor_torch.start_tracking_gradient() @staticmethod def get_description(): return "Constant Tensor constant"