hc/JSCC
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

formatter modified

Browse Source
This commit is contained in:
chun 2024-02-03 20:05:40 +08:00
parent 14cd4ed4e4
commit c81101c269
5 changed files with 13 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
channel.py
View File

@ -1,6 +1,7 @@
import torch import torch
import torch.nn as nn import torch.nn as nn
def channel(channel_type='AWGN', snr=20): def channel(channel_type='AWGN', snr=20):
def AWGN_channel(z_hat: torch.Tensor): def AWGN_channel(z_hat: torch.Tensor):
if z_hat.dim() == 4: if z_hat.dim() == 4:

1
model.py
View File

@ -43,7 +43,6 @@ class _ConvWithPReLU(nn.Module):
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding) self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
self.prelu = nn.PReLU() self.prelu = nn.PReLU()
nn.init.kaiming_normal_(self.conv.weight, mode='fan_out', nonlinearity='leaky_relu') nn.init.kaiming_normal_(self.conv.weight, mode='fan_out', nonlinearity='leaky_relu')
def forward(self, x): def forward(self, x):

3
train.py
View File

@ -19,6 +19,7 @@ from fractions import Fraction
from dataset import Vanilla from dataset import Vanilla
import numpy as np import numpy as np
def set_seed(seed): def set_seed(seed):
np.random.seed(seed) np.random.seed(seed)
torch.manual_seed(seed) torch.manual_seed(seed)
@ -55,7 +56,6 @@ def config_parser():
return parser.parse_args() return parser.parse_args()
def main(): def main():
args = config_parser() args = config_parser()
args.snr_list = list(map(float, args.snr_list)) args.snr_list = list(map(float, args.snr_list))
@ -66,6 +66,7 @@ def main():
for snr in args.snr_list: for snr in args.snr_list:
train(args, ratio, snr) train(args, ratio, snr)
def train(args: config_parser(), ratio: float, snr: float): def train(args: config_parser(), ratio: float, snr: float):
device = torch.device(args.device if torch.cuda.is_available() else 'cpu') device = torch.device(args.device if torch.cuda.is_available() else 'cpu')