foscat 3.0.8__py3-none-any.whl → 3.6.0__py3-none-any.whl

foscat/backend_tens.py

@@ -0,0 +1,63 @@
+import sys
+
+import tensorflow as tf
+
+
+class foscat_backend_tens:
+
+    def __init__(self, backend):
+
+        self.bk = backend
+
+    # ---------------------------------------------−---------
+
+    @tf.function
+    def loss(self, x, batch, loss_function):
+
+        operation = loss_function.scat_operator
+
+        nx = 1
+        if len(x.shape) > 1:
+            nx = x.shape[0]
+
+        with tf.device(
+            operation.gpulist[(operation.gpupos + self.curr_gpu) % operation.ngpu]
+        ):
+            print(
+                "%s Run %d [PROC=%04d] on GPU %s"
+                % (
+                    loss_function.name,
+                    loss_function.id_loss,
+                    self.mpi_rank,
+                    operation.gpulist[
+                        (operation.gpupos + self.curr_gpu) % operation.ngpu
+                    ],
+                )
+            )
+            sys.stdout.flush()
+
+            l_x = x
+            """
+            if nx>1:
+                l_x={}
+            for i in range(nx):
+            """
+
+            if nx == 1:
+                ndata = x.shape[0]
+            else:
+                ndata = x.shape[0] * x.shape[1]
+
+            if self.KEEP_TRACK is not None:
+                l_loss, linfo = loss_function.eval(l_x, batch, return_all=True)
+            else:
+                l_loss = loss_function.eval(l_x, batch)
+
+            g = tf.gradients(l_loss, x)[0]
+            g = self.backend.check_dense(g, ndata)
+            self.curr_gpu = self.curr_gpu + 1
+
+        if self.KEEP_TRACK is not None:
+            return l_loss, g, linfo
+        else:
+            return l_loss, g

foscat/loss_backend_tens.py

@@ -1,49 +1,70 @@
+import sys
+
 import tensorflow as tf
-import numpy as np
+
 
 class loss_backend:
-    
-    def __init__(self,backend):
-        
-        self.bk=backend
+
+    def __init__(self, backend, curr_gpu, mpi_rank):
+
+        self.bk = backend
+        self.curr_gpu = curr_gpu
+        self.mpi_rank = mpi_rank
+
+    def check_dense(self, data, datasz):
+        if isinstance(data, tf.Tensor):
+            return data
+
+        return data.to_dense()
+
     # ---------------------------------------------−---------
-    
+
     @tf.function
-    def loss(self,x,batch,loss_function):
+    def loss(self, x, batch, loss_function, KEEP_TRACK):
+
+        operation = loss_function.scat_operator
 
-        operation=loss_function.scat_operator
+        nx = 1
+        if len(x.shape) > 1:
+            nx = x.shape[0]
 
-        nx=1
-        if len(x.shape)>1:
-            nx=x.shape[0]
-            
-        with tf.device(operation.gpulist[(operation.gpupos+self.curr_gpu)%operation.ngpu]):
-            print('%s Run [PROC=%04d] on GPU %s'%(loss_function.name,self.mpi_rank,
-                                      operation.gpulist[(operation.gpupos+self.curr_gpu)%operation.ngpu]))
+        with tf.device(
+            operation.gpulist[(operation.gpupos + self.curr_gpu) % operation.ngpu]
+        ):
+            print(
+                "%s Run [PROC=%04d] on GPU %s"
+                % (
+                    loss_function.name,
+                    self.mpi_rank,
+                    operation.gpulist[
+                        (operation.gpupos + self.curr_gpu) % operation.ngpu
+                    ],
+                )
+            )
             sys.stdout.flush()
 
-            l_x=x
+            l_x = x
             """
             if nx>1:
                 l_x={}
             for i in range(nx):
             """
-            
-            if nx==1:    
-                ndata=x.shape[0]
+
+            if nx == 1:
+                ndata = x.shape[0]
             else:
-                ndata=x.shape[0]*x.shape[1]
-                
-            if self.KEEP_TRACK is not None:
-                l,linfo=loss_function.eval(l_x,batch,return_all=True)
+                ndata = x.shape[0] * x.shape[1]
+
+            if KEEP_TRACK is not None:
+                l_loss, linfo = loss_function.eval(l_x, batch, return_all=True)
             else:
-                l=loss_function.eval(l_x,batch)
-                
-            g=tf.gradients(l,x)[0]
-            g=self.backend.check_dense(g,ndata)
-            self.curr_gpu=self.curr_gpu+1   
-            
-        if self.KEEP_TRACK is not None:
-            return l,g,linfo
+                l_loss = loss_function.eval(l_x, batch)
+
+            g = tf.gradients(l_loss, x)[0]
+            g = self.check_dense(g, ndata)
+            self.curr_gpu = self.curr_gpu + 1
+
+        if KEEP_TRACK is not None:
+            return l_loss, g, linfo
         else:
-            return l,g
+            return l_loss, g

foscat/loss_backend_torch.py

@@ -1,49 +1,58 @@
-import tensorflow as tf
-import numpy as np
+import torch
+
 
 class loss_backend:
-    
-    def __init__(self,backend):
-        
-        self.bk=backend
+
+    def __init__(self, backend, curr_gpu, mpi_rank):
+
+        self.bk = backend
+        self.curr_gpu = curr_gpu
+        self.mpi_rank = mpi_rank
+
+    def check_dense(self, data, datasz):
+        if isinstance(data, torch.Tensor):
+            return data
+        """
+        idx=tf.cast(data.indices, tf.int32)
+        data=tf.math.bincount(idx,weights=data.values,
+                              minlength=datasz)
+        """
+        return data
+
     # ---------------------------------------------−---------
-    
-    @tf.function
-    def loss(self,x,batch,loss_function):
-
-        operation=loss_function.scat_operator
-
-        nx=1
-        if len(x.shape)>1:
-            nx=x.shape[0]
-            
-        with tf.device(operation.gpulist[(operation.gpupos+self.curr_gpu)%operation.ngpu]):
-            print('%s Run [PROC=%04d] on GPU %s'%(loss_function.name,self.mpi_rank,
-                                      operation.gpulist[(operation.gpupos+self.curr_gpu)%operation.ngpu]))
-            sys.stdout.flush()
-
-            l_x=x
-            """
-            if nx>1:
-                l_x={}
-            for i in range(nx):
-            """
-            
-            if nx==1:    
-                ndata=x.shape[0]
-            else:
-                ndata=x.shape[0]*x.shape[1]
-                
-            if self.KEEP_TRACK is not None:
-                l,linfo=loss_function.eval(l_x,batch,return_all=True)
+
+    def loss(self, x, batch, loss_function, KEEP_TRACK):
+
+        operation = loss_function.scat_operator
+
+        if torch.cuda.is_available():
+            with torch.cuda.device((operation.gpupos + self.curr_gpu) % operation.ngpu):
+
+                l_x = x.clone().detach().requires_grad_(True)
+
+                if KEEP_TRACK is not None:
+                    l_loss, linfo = loss_function.eval(l_x, batch, return_all=True)
+                else:
+                    l_loss = loss_function.eval(l_x, batch)
+
+                l_loss.backward()
+
+                g = l_x.grad
+
+                self.curr_gpu = self.curr_gpu + 1
+        else:
+            l_x = x.clone().detach().requires_grad_(True)
+
+            if KEEP_TRACK is not None:
+                l_loss, linfo = loss_function.eval(l_x, batch, return_all=True)
             else:
-                l=loss_function.eval(l_x,batch)
-                
-            g=tf.gradients(l,x)[0]
-            g=self.backend.check_dense(g,ndata)
-            self.curr_gpu=self.curr_gpu+1   
-            
-        if self.KEEP_TRACK is not None:
-            return l,g,linfo
+                l_loss = loss_function.eval(l_x, batch)
+
+            l_loss.backward()
+
+            g = l_x.grad
+
+        if KEEP_TRACK is not None:
+            return l_loss.detach(), g, linfo
         else:
-            return l,g
+            return l_loss.detach(), g