6 แบบฝึกหัด

class fcf(torch.autograd.Function):
    @staticmethod
    def forward(ctx, zp, w, b):
        # input: zp (N,Mi): (* {\color{blue} $z_j^{(v-1)}$ } *) , w (Mo,Mi): (* {\color{blue} $\partial w^{(v)}_{ji}$ } *) , b (Mo,1): (* {\color{blue} $\partial b^{(v)}_{j}$ } *)
        # output: a (N,Mo): (* {\color{blue} $a_j^{(v)}$ } *)
        zT = torch.transpose(zp, 0, 1) 
        a = w.mm(zT) + b
        ctx.save_for_backward(zp, w, b)
        return torch.transpose(a, 0, 1)
    
    @staticmethod
    def backward(ctx, dEa):
        # input: dEa (N,Mo): (* {\color{blue} $\frac{\partial E}{\partial a^{(v)}_j}$ }*)
        # output: dEzp: (* {\color{blue} $\frac{\partial E}{\partial z^{(v-1)}_i} = \sum_j \frac{\partial E}{\partial a^{(v)}_j} \cdot \frac{\partial a^{(v)}_j}{\partial z^{(v-1)}_i}$ } *) , dEw: (* {\color{blue} $\frac{\partial E}{\partial w^{(v)}_{ji}} = \frac{\partial E}{\partial a^{(v)}_j} \cdot \frac{\partial a^{(v)}_j}{\partial w^{(v)}_{ji}}$ } *) , dEb: (* {\color{blue} $\frac{\partial E}{\partial b^{(v)}_{j}} = \frac{\partial E}{\partial a^{(v)}_j} \cdot \frac{\partial a^{(v)}_j}{\partial b^{(v)}_{j}}$ } *)
        N, _ = dEa.shape
        zp, w, b = ctx.saved_tensors
        dEzp = dEa.mm(w) 
        dEw = torch.transpose(dEa, 0, 1).mm(zp) 
        dEb = torch.transpose(dEa, 0, 1).mm(torch.ones(N,1).to(dEa.device)) 
        return dEzp, dEw, dEb

class MyFCBack(nn.Module):
    def __init__(self, input_channels, num_features):            
        super(MyFCBack, self).__init__()
        self.input_channels = input_channels
        self.num_features = num_features
        sqk = torch.sqrt(torch.Tensor([1/input_channels]))
        initw = 2*sqk*torch.rand(num_features,input_channels)-sqk
        initb = 2*sqk*torch.rand(num_features,1) - sqk
        self.weight = nn.Parameter(initw)
        self.bias = nn.Parameter(initb)
        self.fcf = fcf.apply

    def forward(self, z):
        a = self.fcf(z, self.weight, self.bias)
        return a

class convf(torch.autograd.Function):
    @staticmethod
    def forward(ctx, zp, w, b, S, P):
        '''
        (* {\color{blue} Eq.~\ref{eq: deep conv conv FxCxHxW}:\; 
        $a_{f,k,l}^{(v)} = b_f^{(v)} + \sum_{c=1}^C \sum_{i=1}^{H_F} \sum_{j=1}^{W_F} w_{fcij}^{(v)} \cdot z_{c, S_H \cdot (k-1)+i, S_W \cdot (l-1)+j}^{(v-1)}$ 
        } *)
        (* {\color{blue} zp: $z_{cij}^{(v-1)}$, w: $w_{fcij}^{(v)}$, b: $b_f^{(v)}$, S: stride, P: padding } *)
        '''

        F, C, Hf, Wf = w.shape
        N, D, H, W = zp.shape
        assert C == D, 'Numbers of channels are not matched.'
        
        # Determinte output size
        Ho = int( (H + 2*P - Hf)/S ) + 1
        Wo = int( (W + 2*P - Wf)/S ) + 1

        # Simplify z structure
        simplified_z = MyConv2D._simplify_struct(zp,Hf,Wf,S,P)
        assert simplified_z.shape == (D * Hf * Wf, Ho * Wo * N)
        
        simplified_w = w.view(F,-1)
        assert simplified_w.shape == (F, C * Hf * Wf)
        
        # Compute convolution        
        simplified_out = b + simplified_w.mm(simplified_z)

        # Restructure convoluted output back
        conv_out = simplified_out.view(F, Ho, Wo, N)
        a = conv_out.permute(3, 0, 1, 2) # output (N, M, H', W')

        ctx.save_for_backward(zp, w, b, torch.tensor([S, P]), 
                              simplified_z)
                
        return a       
    
    @staticmethod
    def backward(ctx, dEa):
        # input: dEa (N, F, H', W'): (* {\color{blue} $\delta^{(v)}_{qrs} = \frac{\partial E}{\partial a^{(v)}_{qrs}}$ }*)
        # output: dEzp (N, C, H, W): (* {\color{blue} $\hat{\delta}^{(v-1)}_{fkl} = \frac{\partial E}{\partial z^{(v-1)}_{fkl}} = \sum_{q=1}^F \sum_{r \in \Omega_r} \sum_{s \in \Omega_s} \delta_{qrs}^{(v)} \cdot w_{q,f,k-S_H \cdot (r - 1),l-S_W \cdot (s - 1)}^{(v)}$ } *) 
        #         dEw (F, C, Hf, Wf): (* {\color{blue} $\frac{\partial E}{\partial w_{qfij}^{(v)}} = \sum_{r=1}^{H} \sum_{s=1}^{W} \delta_{qrs}^{(v)} z_{f, S_H \cdot (r-1)+i, S_W \cdot (s-1)+j}^{(v-1)}$ } *)
        #         dEb (F,1): (* {\color{blue} $\frac{\partial E_n}{\partial b_q^{(v)}} = \sum_{r=1}^{H} \sum_{s=1}^{W} \delta_{qrs}^{(v)}$ } *)
        N, F, Ho, Wo = dEa.shape
        zp, w, b, tensorSP, simplified_z = ctx.saved_tensors
        S = tensorSP[0].item()
        P = tensorSP[1].item()
        
        _, C, Hf, Wf = w.shape
        
        # Calculate dEb
        dEb = dEa.sum(dim=(0,2,3)).view(-1,1) # sum over N,H',W'
        
        # Restructure dEa from (N, F, H', W') to (F, H' W' N)        
        simplified_dEa = dEa.permute(1, 2, 3, 0).contiguous().view(F, -1)
        
        # Calculate dEw
        dEw = simplified_dEa.mm(simplified_z.transpose(0,1)) 
        dEw = dEw.view(w.shape) # (F, C, Hf, Wf)
        
        # Calculate dEzp (N, C, H, W): (* {\color{blue} $\hat{\delta}^{(v-1)}_{fkl} = \frac{\partial E}{\partial z^{(v-1)}_{fkl}}$} *)
        # (* {\color{blue} $\frac{\partial E}{\partial z^{(v-1)}_{fkl}} = \sum_{q=1}^F \sum_{r \in \Omega_r} \sum_{s \in \Omega_s} \delta_{qrs}^{(v)} \cdot w_{q,f,k-S_H \cdot (r - 1),l-S_W \cdot (s - 1)}^{(v)} = \sum_{r \in \Omega_r} \sum_{s \in \Omega_s} (\sum_{q=1}^F \delta_{qrs}^{(v)} \cdot w_{q,f,k-S_H \cdot (r - 1),l-S_W \cdot (s - 1)}^{(v)})$ } *) 
        # First, sum over the feature axis
        simplified_w = w.view(F,-1)
        assert simplified_w.shape == (F, C * Hf * Wf)
        
        wdEa_overF = simplified_w.transpose(0,1).mm(
                             simplified_dEa) #(C Hf Wf, H'W'N)
        
        # Sum over spatial indices
        dEzp = convf.sum_omega(wdEa_overF,zp.shape,Hf,Wf,P,S)

        return dEzp, dEw, dEb, None, None
    
    @staticmethod
    def sum_omega(prod_overF, zpshape, Hf, Wf, P, S):
        '''
        Summation over the two omega sets (~over H' and W')
        input: prod_overF (C Hf Wf, H' W' N): (* {\color{blue} $(\sum_{q=1}^F \delta_{qrs}^{(v)} \cdot w_{q,f,k-S_H \cdot (r - 1),l-S_W \cdot (s - 1)}^{(v)})$ } *)
        output: dEzp (N, C, H, W): (* {\color{blue} $\hat{\delta}^{(v-1)}_{fkl} = \frac{\partial E}{\partial z^{(v-1)}_{fkl}} = \sum_{r \in \Omega_r} \sum_{s \in \Omega_s} (\sum_{q=1}^F \delta_{qrs}^{(v)} \cdot w_{q,f,k-S_H \cdot (r - 1),l-S_W \cdot (s - 1)}^{(v)})$ } *) 
        '''
        
        N, C, H, W = zpshape
        H_hat, W_hat = H + 2*P, W + 2*P
        
        # Restructure prod_overF for sum over omega
        prod_overF_reshaped = prod_overF.view(C*Hf*Wf, -1, N)
        prod = prod_overF_reshaped.permute(2, 0, 1).cpu().numpy()
        
        # Prepare result structure
        sum_result = np.zeros((N,C,H_hat,W_hat),dtype=prod.dtype)

        # Get vectorized indices
        c, rx, cx = MyConv2D._get_simplified_indices(zpshape, 
                                                    Hf, Wf, S, P)
        # c.shape = (C Hf Wf, 1)
        # rx.shape = (C Hf Wf, H' W')
        # cx.shape = (C Hf Wf, H' W')
                
        # Sum over omega using np.add.at mechanism
        np.add.at(sum_result, (slice(None), c, rx, cx), prod)
        tsum = torch.tensor(sum_result).to(prod_overF.device)
        
        if P != 0:
            # remove side effect from padding
            return tsum[:, :, P:-P, P:-P]
            
        return tsum

class MyConv2DB(MyConv2D):
    def __init__(self, input_channels, num_kernels, 
                       kernel_size, stride=1, padding=0):            
        super(MyConv2DB, self).__init__(input_channels, 
            num_kernels, kernel_size, stride, padding)
        self.convf = convf.apply
                    
    def forward(self, z):        
        a = self.convf(z, self.weight, self.bias, 
                          self.stride, self.padding)
        return a

class maxpoolf(torch.autograd.Function):
    @staticmethod
    def forward(ctx, zp, Hf=2, Wf=2, S=2, P=0):
        '''
        input: zp (N, C, H, W): (* {\color{gray} $z_{c,i,j}^{(v-1)}$} *)
        output: z (N,C,H',W'): (* {\color{gray} $z_{c,k,l}^{(v)} = g( \{ z_{c, S_H \cdot (k-1)+i, S_W \cdot (l-1)+j}^{(v-1)} \}_{i=1,\ldots, H_F, j=1,\ldots, W_F} )$} *)
        '''

        N, C, H, W = zp.shape
        
        # Determinte output size
        Ho = int( (H + 2*P - Hf)/S ) + 1
        Wo = int( (W + 2*P - Wf)/S ) + 1

        # Restructure zp
        # An operation effect of pooling is different from conv 
        # such that channel/feature axis is treated independently 
        # (like datapoint axis).
                
        restruct_z = zp.view(N * C, 1, H, W)
        sim_z = MyConv2D._simplify_struct(restruct_z, Hf,Wf,S,P)
        assert sim_z.shape == (Hf * Wf, Ho * Wo * N * C)
        
        # Perform pooling function
        # poolz, pool_cache = pool_func(sim_z)
        
        max_idx = torch.argmax(sim_z, dim=0)
        poolz = sim_z[max_idx, range(max_idx.size()[0])]
        pool_cache = max_idx        
        
        # Restructure pooling output
        zpool =  poolz.view(Ho, Wo, N, C)
        z = zpool.permute(2, 3, 0, 1).contiguous()
        
        ctx.save_for_backward(zp, torch.tensor([Hf, Wf, S, P]), 
                              sim_z, pool_cache)
                
        return z
    
    @staticmethod
    def backward(ctx, dEz):
        # input: dEz (N, F, H', W'): (* {\color{blue} $\hat{\delta}_{frs}^{(v)}$} *)
        # output: dEzp (N, F, H, W): (* {\color{blue} $\hat{\delta}_{fkl}^{(v-1)} = \sum_{r \in \Omega_r} \sum_{s \in \Omega_s} \hat{\delta}_{frs}^{(v)} \frac{\partial z_{frs}^{(v)}}{\partial z_{fkl}^{(v-1)}}$ }*)
        zp, tensorHfWfSP, sim_z, pool_cache = ctx.saved_tensors
        Hf = tensorHfWfSP[0].item()
        Wf = tensorHfWfSP[1].item()
        S = tensorHfWfSP[2].item()
        P = tensorHfWfSP[3].item()
        
        N, F, H, W = zp.shape
    
        sim_dEa = torch.zeros(sim_z.shape).to(zp.device) 
        sim_dEz = dEz.permute(2, 3, 0, 1).contiguous().view(-1,) 
        
        # Perform dpooling function
        # sim_dEa = dpool_func(poolz, pool_cache)
        # sim_dEa:  (* {\color{blue} $\hat{\delta}_{frs}^{(v)} \cdot \frac{\partial z_{frs}^{(v)}}{\partial z_{fkl}^{(v-1)}} = \left\{\begin{array}{l l}\delta_{frs}^{(v)} & \;\mbox{when}\; f, k, l \;\mbox{are the max id's}\\ 0 & \;\mbox{otherwise}\end{array}\right.$} *)
        
        sim_dEa[pool_cache, range(pool_cache.size()[0])] = 
            sim_dEz    # (Hf Wf, H' W' N F)
                
        # Sum over spatial indices
        dEzp_sim = convf.sum_omega(sim_dEa, 
                                   (N*F, 1, H, W), Hf, Wf, P, S)
        dEzp = dEzp_sim.view(zp.shape)

        return dEzp, None, None, None, None
    

class MyMaxpool(nn.Module):
    def __init__(self, kernel_size, stride, padding=0):            
        super(MyMaxpool, self).__init__()
        self.maxpoolf = maxpoolf.apply
        self.Hf = kernel_size
        self.Wf = kernel_size
        self.stride = stride
        self.padding = padding  
                            
    def forward(self, zp):   
        z = self.maxpoolf(zp, self.Hf, self.Wf, 
                          self.stride, self.padding)
        return z