Verified Commit c76ec438 authored by Max Ehrlich's avatar Max Ehrlich
Browse files

Add proper exports

parent df54b73f
from .codec import *
from .tensors import *
\ No newline at end of file
......@@ -6,7 +6,7 @@ C = torch.einsum('ijab,abg,gk->ijk', (D(), Z(), S()))
C_i = torch.einsum('ijab,abg,gk->ijk', (D(), Z(), S_i()))
def codec(image_size, block_size):
def codec(image_size, block_size=(8, 8)):
B_i = B(image_size, block_size)
J = torch.einsum('srxyij,ijk->srxyk', (B_i, C))
J_i = torch.einsum('srxyij,ijk->xyksr', (B_i, C_i))
......
from .avgpool import *
from .batchnorm import *
from .convolution import *
from .relu import *
\ No newline at end of file
import torch
class JpegAvgPool(torch.nn.modules.Module):
class AvgPool(torch.nn.modules.Module):
def __init__(self):
super(JpegAvgPool, self).__init__()
super(AvgPool, self).__init__()
def forward(self, input):
result = torch.mean(input[:, :, :, :, 0].view(-1, input.shape[1], input.shape[2]*input.shape[3]), 2)
......
import torch
class JpegBatchNorm(torch.nn.modules.Module):
class BatchNorm(torch.nn.modules.Module):
def __init__(self, bn):
super(JpegBatchNorm, self).__init__()
super(BatchNorm, self).__init__()
self.mean = bn.running_mean
self.var = bn.running_var
......
import torch
import opt_einsum as oe
from jpeg_codec.tensors import D_n, D, Z, S_i, S
from jpeg_codec import D_n, D, Z, S_i, S
class JpegRelu(torch.nn.modules.Module):
class ReLU(torch.nn.modules.Module):
def __init__(self, n_freqs):
super(JpegRelu, self).__init__()
super(ReLU, self).__init__()
self.C_n = torch.einsum('ijab,abg,gk->ijk', [D_n(n_freqs), Z(), S_i()])
self.Hm = torch.einsum('ijab,ijuv,abg,gk,uvh,hl->ijkl', [D(), D(), Z(), S_i(), Z(), S()])
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment