diff --git a/lib/models/include/models/candle_uno/candle_uno.h b/lib/models/include/models/candle_uno/candle_uno.h index a2d21f2830..e3093972d7 100644 --- a/lib/models/include/models/candle_uno/candle_uno.h +++ b/lib/models/include/models/candle_uno/candle_uno.h @@ -32,7 +32,7 @@ CandleUnoConfig get_default_candle_uno_config(); * this model. * * @param CandleUnoConfig The config of the Candle Uno model. - * @return ComputationGraph The PCG of a Transformer model. + * @return ComputationGraph The computation graph of a Candle Uno model. */ ComputationGraph get_candle_uno_computation_graph(CandleUnoConfig const &); diff --git a/lib/models/include/models/simvp/simvp.h b/lib/models/include/models/simvp/simvp.h new file mode 100644 index 0000000000..8485ed58cf --- /dev/null +++ b/lib/models/include/models/simvp/simvp.h @@ -0,0 +1,75 @@ +#ifndef _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_SIMVP_H +#define _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_SIMVP_H + +#include "pcg/computation_graph_builder.h" +#include "simvp_config.dtg.h" + +namespace FlexFlow { + +// Helper functions to construct the SimVP model + +/** + * @brief Get the default configs of SimVP model. + */ +SimVPConfig get_default_simvp_config(); + +std::vector create_simvp_samplings(size_t N_S, bool reverse = false); + +tensor_guid_t create_simvp_convsc(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input, + size_t in_dim, + size_t out_dim, + int kernel_size = 3, + bool downsampling = false, + bool upsampling = false); + +tensor_guid_t create_simvp_gsta_meta_block(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input, + int in_channels, + int out_channels, + float mlp_ratio = 8.0, + float drop = 0.0, + float drop_path = 0.0); + +tensor_guid_t create_simvp_ga_sub_block(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input, + int dim, + int kernel_size = 21, + float mlp_ratio = 4.0, + float drop = 0.0, + float drop_path = 0.1, + float init_value = 1e-2); + +std::pair + create_simvp_encoder(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input); + +tensor_guid_t create_simvp_middle_net(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &embed, + int channel_in, + int channel_hid, + float mlp_ratio = 4.0, + float drop = 0.0, + float drop_path = 0.1); + +tensor_guid_t create_simvp_decoder(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &hid, + tensor_guid_t const &skip); + +/** + * @brief Get the SimVP computation graph. + * + * @param SimVPConfig The config of the SimVP model. + * @return ComputationGraph The computation graph of a SimVP model. + */ +ComputationGraph get_simvp_computation_graph(SimVPConfig const &config); + +} // namespace FlexFlow + +#endif diff --git a/lib/models/include/models/simvp/simvp_config.struct.toml b/lib/models/include/models/simvp/simvp_config.struct.toml new file mode 100644 index 0000000000..a753f1f237 --- /dev/null +++ b/lib/models/include/models/simvp/simvp_config.struct.toml @@ -0,0 +1,72 @@ +namespace = "FlexFlow" +name = "SimVPConfig" + +features = [ + "eq", + "ord", + "hash", + "json", + "rapidcheck", + "fmt", +] + +includes = [ + "", + "", + "", +] + +src_includes = [ + "utils/fmt/vector.h", + "utils/fmt/map.h", + "utils/hash/vector.h", + "utils/hash/map.h", +] + +[[fields]] +name = "batch_size" +type = "size_t" + +[[fields]] +name = "hid_S" +type = "size_t" + +[[fields]] +name = "hid_T" +type = "size_t" + +[[fields]] +name = "N_S" +type = "size_t" + +[[fields]] +name = "N_T" +type = "size_t" + +[[fields]] +name = "model_type" +type = "std::string" + +[[fields]] +name = "mlp_ratio" +type = "float" + +[[fields]] +name = "drop" +type = "float" + +[[fields]] +name = "drop_path" +type = "float" + +[[fields]] +name = "spatio_kernel_enc" +type = "size_t" + +[[fields]] +name = "spatio_kernel_dec" +type = "size_t" + +[[fields]] +name = "in_shape" +type = "std::vector" diff --git a/lib/models/src/models/simvp/simvp.cc b/lib/models/src/models/simvp/simvp.cc new file mode 100644 index 0000000000..da2753cc3f --- /dev/null +++ b/lib/models/src/models/simvp/simvp.cc @@ -0,0 +1,250 @@ +#include "models/simvp/simvp.h" + +namespace FlexFlow { + +SimVPConfig get_default_simvp_config() { + return SimVPConfig{ + /*batch_size=*/1, + /*hid_S=*/16, + /*hid_T=*/256, + /*N_S=*/4, + /*N_T=*/4, + /*model_type=*/"gSTA", + /*mlp_ratio=*/8.0, + /*drop=*/0.0, + /*drop_path=*/0.0, + /*spatio_kernel_enc=*/3, + /*spatio_kernel_dec=*/3, + /*in_shape=*/ + {10, 3, 32, 32}, + }; +} + +std::vector create_simvp_samplings(size_t N_S, bool reverse) { + size_t N_S_even_floor = (N_S / 2) * 2; + + auto const change_to_true = [&](size_t idx) -> bool { + return (reverse == false) ? (idx % 2 == 1) : (idx % 2 == 0); + }; + + std::vector samplings(N_S_even_floor, false); + for (size_t i = 0; i < N_S_even_floor; i++) { + if (change_to_true(i)) { + samplings[i] = true; + } + } + + return samplings; +} + +tensor_guid_t create_simvp_convsc(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input, + size_t in_dim, + size_t out_dim, + int kernel_size, + bool downsampling, + bool upsampling) { + int stride = (downsampling == true) ? 2 : 1; + int padding = (kernel_size - stride + 1) / 2; + int out_channels = upsampling ? out_dim * 4 : out_dim; + + tensor_guid_t conv_out = cgb.conv2d(input, + out_channels, + kernel_size, + kernel_size, + stride, + stride, + padding, + padding); + + return conv_out; +} + +std::pair + create_simvp_encoder(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input) { + size_t C = config.in_shape.at(1); // Channel + std::vector samplings = create_simvp_samplings(config.N_S); + + tensor_guid_t enc1 = create_simvp_convsc(cgb, + config, + input, + C, + config.hid_S, + config.spatio_kernel_enc, + samplings[0], + false); + tensor_guid_t latent = enc1; + + for (size_t i = 1; i < samplings.size(); i++) { + latent = create_simvp_convsc(cgb, + config, + latent, + config.hid_S, + config.hid_S, + config.spatio_kernel_enc, + samplings[i], + false); + } + + return {latent, enc1}; +} + +// TODO +tensor_guid_t create_simvp_ga_sub_block(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input, + int dim, + int kernel_size, + float mlp_ratio, + float drop, + float drop_path, + float init_value) { + return input; +} + +tensor_guid_t create_simvp_gsta_meta_block(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input, + int in_channels, + int out_channels, + float mlp_ratio, + float drop, + float drop_path) { + tensor_guid_t z = create_simvp_ga_sub_block( + cgb, config, input, in_channels, 21, mlp_ratio, drop, drop_path); + + if (in_channels == out_channels) { + return z; + } else { + return cgb.conv2d(z, out_channels, 1, 1, 1, 1, 0, 0); + } +} + +tensor_guid_t create_simvp_middle_net(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &embed, + int channel_in, + int channel_hid, + float mlp_ratio, + float drop, + float drop_path) { + if (config.model_type != "gSTA") { + throw mk_runtime_error( + fmt::format("Currently only model_type=gSTA is " + "supported, but found model_type={}. " + "If you need support for additional " + "model_type values, please create an issue.", + config.model_type)); + } + + tensor_guid_t z = embed; + + // Downsample + z = create_simvp_gsta_meta_block( + cgb, config, z, channel_in, channel_hid, mlp_ratio, drop, drop_path); + + // Middle layers + for (size_t i = 1; i < config.N_T - 1; i++) { + z = create_simvp_gsta_meta_block( + cgb, config, z, channel_hid, channel_hid, mlp_ratio, drop, drop_path); + } + + // Upsample + z = create_simvp_gsta_meta_block( + cgb, config, z, channel_hid, channel_in, mlp_ratio, drop, drop_path); + + return z; +} + +tensor_guid_t create_simvp_decoder(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &hid, + tensor_guid_t const &skip) { + + std::cout << "hid shape: " << cgb.get_shape(hid) << std::endl; + std::cout << "skip shape: " << cgb.get_shape(skip) << std::endl; + + size_t C = config.in_shape.at(1); // Channel + std::vector samplings = create_simvp_samplings(config.N_S, true); + + tensor_guid_t latent = hid; + for (size_t i = 0; i < samplings.size() - 1; i++) { + latent = create_simvp_convsc(cgb, + config, + latent, + config.hid_S, + config.hid_S, + config.spatio_kernel_dec, + false, + samplings[i]); + } + + tensor_guid_t out = create_simvp_convsc(cgb, + config, + cgb.add(latent, skip), + config.hid_S, + config.hid_S, + config.spatio_kernel_dec, + false, + samplings[samplings.size() - 1]); + + return cgb.conv2d(out, C, 1, 1, 1, 1, 0, 0, std::nullopt); +} + +ComputationGraph get_simvp_computation_graph(SimVPConfig const &config) { + ComputationGraphBuilder cgb; + + // Create input tensor + size_t B = config.batch_size; // Number of samples + size_t T = config.in_shape.at(0); // Number of frames in each sample + size_t C = config.in_shape.at(1); // Channel + size_t H = config.in_shape.at(2); // Height + size_t W = config.in_shape.at(3); // Width + + // std::cout << "B T C H W: " << B << " " << T << " " << C << " " << H << " " + // << W << std::endl; + + TensorShape input_shape = TensorShape{ + TensorDims{FFOrdered{B * T, C, H, W}}, + DataType::FLOAT, + }; + tensor_guid_t input = cgb.create_input(input_shape, CreateGrad::YES); + + // Create the model + auto [embed, skip] = create_simvp_encoder(cgb, config, input); + + // std::cout << "embed shape: " << cgb.get_shape(embed) << std::endl; + + // TODO: need to reshape embed here + + tensor_guid_t hid = create_simvp_middle_net(cgb, + config, + embed, + T * config.hid_S, + config.hid_T, + config.mlp_ratio, + config.drop, + config.drop_path); + + // TODO: need to reshape hid here + // std::cout << "hid shape: " << cgb.get_shape(hid) << std::endl; + // auto const embed_shape = cgb.get_shape(embed).dims.ff_ordered; + // std::vector hid_shape = { + // static_cast(B * T), + // static_cast(embed_shape.at(ff_dim_t{1})), + // static_cast(embed_shape.at(ff_dim_t{2})), + // static_cast(embed_shape.at(ff_dim_t{3})), + // }; + // hid = cgb.reshape(hid, hid_shape); + // std::cout << "hid new shape: " << cgb.get_shape(hid) << std::endl; + + // TODO: enable below after reshaping + // tensor_guid_t output = create_simvp_decoder(cgb, config, hid, skip); + + return cgb.computation_graph; +} + +} // namespace FlexFlow diff --git a/lib/models/test/src/models/simvp/simvp.cc b/lib/models/test/src/models/simvp/simvp.cc new file mode 100644 index 0000000000..c6939b7d20 --- /dev/null +++ b/lib/models/test/src/models/simvp/simvp.cc @@ -0,0 +1,19 @@ +#include "models/simvp/simvp.h" +#include "pcg/computation_graph.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("get_simvp_computation_graph") { + SimVPConfig config = get_default_simvp_config(); + + ComputationGraph result = get_simvp_computation_graph(config); + + SUBCASE("num layers") { + int result_num_layers = get_layers(result).size(); + int correct_num_layers = 19; + CHECK(result_num_layers == correct_num_layers); + } + } +}