I want to rebuild a customized UNET in tensorflow.js, but I struggle to concatenate layers like described in the decoder block class:
class decoder_block(nn.Module):
def __init__(self, in_c, out_c):
super().__init__()
self.up = nn.ConvTranspose2d(in_c, out_c, kernel_size=2, stride=2, padding=0)
self.conv = conv_block(out_c+out_c, out_c)
def forward(self, inputs, skip):
x = self.up(inputs)
x = torch.cat([x, skip], axis=1)
x = self.conv(x)
return x
from https://medium.com/analytics-vidhya/unet-implementation-in-pytorch-idiot-developer-da40d955f201
However, in tensorflow.js it seems not to be straightforward. For testing I tried to concatenate a layer with itself to have not any kind of dimension issues:
model.add(tf.layers.conv2d({
kernelSize: [3, 3],
filters: 64,
strides: 1,
activation: 'relu',
kernelInitializer: 'randomUniform',
name:"s1"
}));
console.log(model.getLayer("s1").getWeights())
const concatLayer = tf.layers.concatenate();
const output = concatLayer.apply(
model.getLayer("s1").getWeights()
,model.getLayer("s1").getWeights());
But I got the following error:
var _this = _super.call(this, message) || this;
^
ValueError: A `Concatenate` layer requires inputs with matching shapes except for the concat axis. Got input shapes: [[3,3,64,64],[64]]
So is there a way to do the above in tensorflow.js?