I was trying to build a very simple model to create simple circle images with 5 different colors and random radius, but I did not get a good result. I have created 500 images for each color for the training set.
I am just wondering that the logic of the model is correct.
The input shape is :
one_hot([1,0,0,0,0]) and the output shape is
Here is my simple model:
cls_model( (model): Sequential( (0): Linear(in_features=5, out_features=1024, bias=True) (1): ReLU() (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True) (3): Linear(in_features=1024, out_features=961, bias=True) (4): ReLU() ) (tconv1): ConvTranspose2d(1, 1, kernel_size=(3, 3), stride=(1, 1)) (tconv2): ConvTranspose2d(1, 3, kernel_size=(3, 3), stride=(1, 1)) )
criterion = nn.MSELoss() optim = torch.optim.Adam(net.parameters(), lr=1e-3)