ChannelCoding_AE/decoders.py at main · Yitiann/ChannelCoding_AE · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

import torch.nn.functional as F
from torch.nn import Parameter
import torch

from numpy import arange
from numpy.random import mtrand
import numpy as np


#####################
# RNN decoder
####################

class RNN_decoder_rate2(torch.nn.Module):
    def __init__(self, args):
        super(RNN_decoder_rate2, self).__init__()
        self.args = args

        use_cuda = not args.no_cuda and torch.cuda.is_available()
        self.this_device = torch.device("cuda" if use_cuda else "cpu")

        self.dec_linear = torch.nn.Linear(args.code_rate_n * args.block_len, args.code_rate_n * args.block_len)
        self.dec_rnn = torch.nn.GRU(args.code_rate_n, args.dec_num_unit,
                                    num_layers=2, bias=True, batch_first=True,
                                    dropout=0, bidirectional=True)

        self.final = torch.nn.Linear(2 * args.dec_num_unit, 1)

    def set_parallel(self):
        pass

    def forward(self, received):

        received = received.type(torch.FloatTensor).to(self.this_device) #(batch_size, n*block_len,1)
        permuted = received.permute(0, 2, 1) #(batch_size, 1, n*block_len)
        code = self.dec_linear(permuted) #(batch_size, 1, n*block_len)
        code= code.view(self.args.batch_size,self.args.block_len,self.args.code_rate_n) #(batch_size,block_len,n) [ab,ab]

        #  Decoder
        x_plr  = self.dec_rnn(code)[0]
        final  = torch.sigmoid(self.final(x_plr))
        return final

class RNN_decoder_rate_high(torch.nn.Module):
    def __init__(self, args):
        super(RNN_decoder_rate_high, self).__init__()
        self.args = args

        use_cuda = not args.no_cuda and torch.cuda.is_available()
        self.this_device = torch.device("cuda" if use_cuda else "cpu")
        self.dec_linear = torch.nn.Linear(args.code_rate_n * args.block_len, 2* args.code_rate_k * args.block_len)

        self.dec_rnn = torch.nn.GRU(args.code_rate_n, args.dec_num_unit,
                                    num_layers=2, bias=True, batch_first=True,
                                    dropout=0, bidirectional=True)

        self.final = torch.nn.Linear(2 * args.dec_num_unit, args.code_rate_k)

    def set_parallel(self):
        pass

    def forward(self, received):

        received = received.type(torch.FloatTensor).to(self.this_device) #(batch_size, n*block_len,1)
        permuted = received.permute(0, 2, 1) #(batch_size, 1, n*block_len)
        code = self.dec_linear(permuted) #(batch_size, 1, n*block_len)
        code= code.view(self.args.batch_size,self.args.block_len, self.args.code_rate_n) #(batch_size,block_len,n) [ab,ab]

        #  Decoder
        x_plr      = self.dec_rnn(code)[0]
        final      = torch.sigmoid(self.final(x_plr))
        return final