foscat 3.1.6__py3-none-any.whl → 3.2.0__py3-none-any.whl

foscat/backend_tens.py

@@ -1,49 +1,62 @@
+import sys
+
 import tensorflow as tf
-import numpy as np
+
 
 class foscat_backend_tens:
-    
-    def __init__(self,backend):
-        
-        self.bk=backend
+
+    def __init__(self, backend):
+
+        self.bk = backend
+
     # ---------------------------------------------−---------
-    
+
     @tf.function
-    def loss(self,x,batch,loss_function):
+    def loss(self, x, batch, loss_function):
 
-        operation=loss_function.scat_operator
+        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]))
+        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
+            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]
-                
+                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)
+                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   
-            
+                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,g,linfo
+            return l_loss, g, linfo
         else:
-            return l,g
+            return l_loss, g

foscat/loss_backend_tens.py

@@ -1,60 +1,70 @@
-import tensorflow as tf
-import numpy as np
 import sys
 
+import tensorflow as tf
+
+
 class loss_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):
+
+    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,KEEP_TRACK):
+    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]
-            
-        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]))
+        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
+            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]
-                
+                ndata = x.shape[0] * x.shape[1]
+
             if KEEP_TRACK is not None:
-                l,linfo=loss_function.eval(l_x,batch,return_all=True)
+                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.check_dense(g,ndata)
-            self.curr_gpu=self.curr_gpu+1   
-            
+                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,g,linfo
+            return l_loss, g, linfo
         else:
-            return l,g
+            return l_loss, g

foscat/loss_backend_torch.py

@@ -1,18 +1,15 @@
 import torch
-from torch.autograd import grad
-import numpy as np
-import sys
+
 
 class loss_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):
+
+    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
         """
@@ -21,64 +18,42 @@ class loss_backend:
                               minlength=datasz)
         """
         return data
-        
-    # ---------------------------------------------−---------
 
-    def loss(self,x,batch,loss_function,KEEP_TRACK):
+    # ---------------------------------------------−---------
 
-        operation=loss_function.scat_operator
+    def loss(self, x, batch, loss_function, KEEP_TRACK):
 
-        nx=1
-        if len(x.shape)>1:
-            nx=x.shape[0]
+        operation = loss_function.scat_operator
 
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         if torch.cuda.is_available():
-            with torch.cuda.device((operation.gpupos+self.curr_gpu)%operation.ngpu):
+            with torch.cuda.device((operation.gpupos + self.curr_gpu) % operation.ngpu):
+
+                l_x = x.clone().detach().requires_grad_(True)
 
-                l_x=x.clone().detach().requires_grad_(True)
-            
-                if nx==1:    
-                    ndata=x.shape[0]
-                else:
-                    ndata=x.shape[0]*x.shape[1]
-                
                 if KEEP_TRACK is not None:
-                    l,linfo=loss_function.eval(l_x,batch,return_all=True)
+                    l_loss, linfo = loss_function.eval(l_x, batch, return_all=True)
                 else:
-                    l=loss_function.eval(l_x,batch)
+                    l_loss = loss_function.eval(l_x, batch)
 
-                l.backward()
-            
-                g=l_x.grad
-            
-                self.curr_gpu=self.curr_gpu+1
+                l_loss.backward()
+
+                g = l_x.grad
+
+                self.curr_gpu = self.curr_gpu + 1
         else:
-            l_x=x.clone().detach().requires_grad_(True)
+            l_x = x.clone().detach().requires_grad_(True)
 
-            if nx==1:    
-                ndata=x.shape[0]
-            else:
-                ndata=x.shape[0]*x.shape[1]
 
             if KEEP_TRACK is not None:
-                l,linfo=loss_function.eval(l_x,batch,return_all=True)
+                l_loss, linfo = loss_function.eval(l_x, batch, return_all=True)
             else:
-                """
-                sx=operation.eval(l_x)
-                tmp=(sx.C01-1.0)
-    
-                l=operation.backend.bk_reduce_sum(tmp*tmp) #loss_function.eval(l_x,batch)
-                """
-                
-                l=loss_function.eval(l_x,batch)
-
-            l.backward()
-            
-            g=l_x.grad
-            
+                l_loss = loss_function.eval(l_x, batch)
+
+            l_loss.backward()
+
+            g = l_x.grad
+
         if KEEP_TRACK is not None:
-            return l.detach(),g,linfo
+            return l_loss.detach(), g, linfo
         else:
-            return l.detach(),g
-
+            return l_loss.detach(), g