Forcing PyTorch Neural Net to output a specific datatype

Question:

I am learning how to create a GAN with PyTorch 1.12 and I need the instance returned by my generator to fall into a specific feature space.

The model in my generator class looks like this:

self.model = nn.Sequential(
            nn.Linear(2, 16),
            nn.ReLU(),
            nn.Linear(16, 32),
            nn.ReLU(),
            nn.Linear(32, 2),
        )

I need every feature in the instance returned by my generator to be an unsigned integer. The noise fed into the generator fits this, but the output has decimals and some features are negative. How can I set this?

(If needed I can include more of my code but I figured this is the relevant line and wanted to keep this question short to make things easier for everyone.)

Asked By: pantman

||

Answers:

You can try clipping negative values and casting to torch.int32:

from torch import nn
import torch


class TransformOutput(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return torch.clamp_min(x, min=0).to(dtype=torch.int32)


model = nn.Sequential(
    nn.Linear(2, 16),
    nn.ReLU(),
    nn.Linear(16, 32),
    nn.ReLU(),
    nn.Linear(32, 2),
    TransformOutput(),
)

x = torch.rand(4, 2) - 0.5
out = model(x)
assert out.dtype == torch.int32
Answered By: u1234x1234